From 89989afb9a2781d4f5206267feb7bcc8a0abaa67 Mon Sep 17 00:00:00 2001 From: yangkaixin Date: Wed, 26 Aug 2026 16:10:40 +0800 Subject: [PATCH 1/5] auto scope is working on solve_tril, and offer 12% performence improvment over simt_only now --- docs/zh/examples/08_costmodel_example.md | 111 + third_party/ascend/CMakeLists.txt | 1 + third_party/ascend/ascend_ir.cc | 22 +- third_party/ascend/backend/compiler.py | 293 +- third_party/ascend/backend/utils.py | 7 +- .../AscendModel/Analysis/HIVMAnalysis.h | 11 +- .../SimtAnchorAnalysis.h | 82 +- .../StagePartitioner.h | 28 +- .../RouteModel/SimdSimtCostModel.h | 248 +- .../AscendModel/RouteModel/StageCostModels.h | 65 - .../RouteModel/StageRouteCostModel.h | 34 +- .../include/AscendModel/Transforms/Passes.td | 12 +- .../SimtSelection.h | 0 .../lib/AscendModel/Analysis/HIVMAnalysis.cpp | 554 +--- .../SimtAnchorAnalysis.cpp | 331 +-- .../StagePartitioner.cpp | 1232 +++------ .../lib/AscendModel/RouteModel/CMakeLists.txt | 8 +- .../RouteModel/SimdSimtCostModel.cpp | 2405 ++--------------- .../RouteModel/StageCostModels.cpp | 686 +---- .../RouteModel/StageRouteCostModel.cpp | 412 ++- .../lib/AscendModel/Transforms/CMakeLists.txt | 1 + .../Transforms/HIVMAnalysisPass.cpp | 17 +- .../Transforms/MaterializeSimtScopes.cpp | 35 +- .../Transforms/SelectSimdSimtCostModel.cpp | 137 +- .../ascend/costmodel/profiles/README.md | 2 +- .../simd_simt/david_v100_des_feedback_v1.json | 18 - .../simd_simt/david_v100_simd_simt_v1.json | 8 +- third_party/ascend/include/CMakeLists.txt | 3 +- .../include/TTIRLayoutMerge/CMakeLists.txt | 3 + .../ascend/include/TTIRLayoutMerge/Passes.h | 21 + .../ascend/include/TTIRLayoutMerge/Passes.td | 18 + .../RowCoalescing.h | 6 +- .../TTIRLayoutMergePass.h | 2 +- .../ascend/include/TritonToLinalg/Passes.h | 1 - .../ascend/include/TritonToLinalg/Passes.td | 12 - third_party/ascend/lib/CMakeLists.txt | 2 + .../DiscreteMaskAccessConversionPass.cpp | 2 +- .../ascend/lib/TTIRLayoutMerge/CMakeLists.txt | 18 + .../RowCoalescing.cpp | 3 +- .../TTIRLayoutMergePass.cpp | 4 +- .../ascend/lib/TritonToLinalg/CMakeLists.txt | 2 - .../StridedLoadStoreRewrite.cpp | 2 +- .../lib/TritonToLinalg/TritonOpConverter.cpp | 2 +- .../lib/TritonToLinalg/TritonToLinalgPass.cpp | 33 +- .../UnstructureConversionPass.cpp | 2 +- third_party/ascend/triton_ascend.cc | 69 +- .../General/TritonToLinalg/simt_scope.mlir | 6 + .../unittest/costmodel_ut/PassesTest.cpp | 313 ++- .../costmodel_ut/SimdSimtCostModelTest.cpp | 843 +++--- .../test_compiler_costmodel_contract.py | 32 +- .../pytest_ut/test_auto_blockify_v1_policy.py | 96 +- .../test_simd_simt_costmodel_cases.py | 440 +++ 52 files changed, 2846 insertions(+), 5849 deletions(-) create mode 100644 docs/zh/examples/08_costmodel_example.md rename third_party/ascend/costmodel/include/AscendModel/{RouteModel => Analysis}/SimtAnchorAnalysis.h (64%) rename third_party/ascend/costmodel/include/AscendModel/{RouteModel => Analysis}/StagePartitioner.h (80%) rename third_party/ascend/costmodel/include/AscendModel/{RouteModel => Transforms}/SimtSelection.h (100%) rename third_party/ascend/costmodel/lib/AscendModel/{RouteModel => Analysis}/SimtAnchorAnalysis.cpp (68%) rename third_party/ascend/costmodel/lib/AscendModel/{RouteModel => Analysis}/StagePartitioner.cpp (53%) rename third_party/ascend/costmodel/lib/AscendModel/{RouteModel => }/Transforms/MaterializeSimtScopes.cpp (85%) rename third_party/ascend/costmodel/lib/AscendModel/{RouteModel => }/Transforms/SelectSimdSimtCostModel.cpp (62%) delete mode 100644 third_party/ascend/costmodel/profiles/simd_simt/david_v100_des_feedback_v1.json create mode 100644 third_party/ascend/include/TTIRLayoutMerge/CMakeLists.txt create mode 100644 third_party/ascend/include/TTIRLayoutMerge/Passes.h create mode 100644 third_party/ascend/include/TTIRLayoutMerge/Passes.td rename third_party/ascend/include/{TritonToLinalg => TTIRLayoutMerge}/RowCoalescing.h (92%) rename third_party/ascend/include/{TritonToLinalg => TTIRLayoutMerge}/TTIRLayoutMergePass.h (93%) create mode 100644 third_party/ascend/lib/TTIRLayoutMerge/CMakeLists.txt rename third_party/ascend/lib/{TritonToLinalg => TTIRLayoutMerge}/RowCoalescing.cpp (99%) rename third_party/ascend/lib/{TritonToLinalg => TTIRLayoutMerge}/TTIRLayoutMergePass.cpp (96%) create mode 100644 third_party/ascend/unittest/pytest_ut/test_simd_simt_costmodel_cases.py diff --git a/docs/zh/examples/08_costmodel_example.md b/docs/zh/examples/08_costmodel_example.md new file mode 100644 index 0000000000..adc345bbea --- /dev/null +++ b/docs/zh/examples/08_costmodel_example.md @@ -0,0 +1,111 @@ +# Costmodel 端到端示例 + +本示例展示 costmodel 后端的基本调用流程： + +- 使用 Triton 前端算子生成 TTIR； +- 为多个候选 config 构造 `costmodel_bench` 输入； +- 调用 `costmodel_bench` 得到每个 config 的预测耗时。 + +这个流程适合在 autotune 前快速筛掉预计性能较差的 config。示例只使用向量加法 kernel，便于聚焦 costmodel 的输入和返回值。 + +## 完整示例 + +将下面的代码保存为 `costmodel_example.py` 后运行： + +```python +from __future__ import annotations + +import triton +import triton.language as tl +from triton.backends.ascend.runtime.costmodel_runtime import costmodel_bench +from triton.backends.compiler import GPUTarget +from triton.compiler import ASTSource +from triton.compiler.code_generator import ast_to_ttir +from triton.compiler.compiler import make_backend +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir + + +@triton.jit +def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + tl.store(output_ptr + offsets, x + y, mask=mask) + + +def make_ttir(kernel, signature, constants): + source = ASTSource(kernel, signature, constants, attrs=None) + target = GPUTarget("npu", "", 32) + backend = make_backend(target) + + options = backend.parse_options( + { + "num_warps": 8, + "num_stages": 2, + "debug": False, + "multibuffer": False, + "compile_mode": "simd", + "enable_costmodel_backend": True, + **source.parse_options(), + } + ) + + context = ir.context() + ir.load_dialects(context) + ascend_ir.load_dialects(context) + return str(ast_to_ttir(kernel, source, context, options, {}, {})) + + +signature = { + "x_ptr": "*fp32", + "y_ptr": "*fp32", + "output_ptr": "*fp32", + "n_elements": "i32", +} +n_elements = 98432 +configs = [ + {"name": "block256", "BLOCK_SIZE": 256}, + {"name": "block1024", "BLOCK_SIZE": 1024}, + {"name": "block2048", "BLOCK_SIZE": 2048}, +] + +items = [] +for cfg in configs: + ttir = make_ttir(add_kernel, signature, {"BLOCK_SIZE": cfg["BLOCK_SIZE"]}) + items.append( + { + "config": cfg["name"], + "ttir": ttir, + # n_elements 是 signature 中的第 4 个参数，对应 TTIR 里的 %arg3。 + # pid_x 给 tl.program_id(0) 一个静态估算值。 + "arg_bindings": f"arg3={n_elements},pid_x=0", + } + ) + +latencies = costmodel_bench(items) +for config, latency_us in sorted(latencies.items(), key=lambda item: item[1]): + print(f"{config}: {latency_us:.3f} us") +``` + +## 示例输出 + +不同版本的 costmodel 参数可能会使具体数值略有不同，但输出结构类似： + +```text +block256: 0.098 us +block1024: 0.110 us +block2048: 0.126 us +``` + +`costmodel_bench` 的返回值是一个字典，key 为传入的 `config`，value 为预测耗时，单位是微秒。上层 autotune 逻辑可以按 value 排序，优先保留预测更快的 config。 + +## 关键点说明 + +1. `ASTSource + ast_to_ttir` 只生成 TTIR，不会真实编译或启动 kernel。 +2. `config` 会影响 `tl.constexpr`，例如 `BLOCK_SIZE`，因此每个候选 config 都需要生成各自的 TTIR。 +3. `costmodel_bench` 接收的每个元素至少包含 `config` 和 `ttir`，也可以附带 `arg_bindings`。 +4. `arg_bindings` 用于把运行时整数参数绑定到 TTIR 中的 `%argN`。例如本例中 `n_elements=98432` 对应 `arg3=98432`。 +5. 如果 kernel 中使用 `tl.program_id(0)`，通常需要传入 `pid_x=0`。如果还使用 `tl.num_programs(0)`，可额外传入 `num_programs_x=...`。 diff --git a/third_party/ascend/CMakeLists.txt b/third_party/ascend/CMakeLists.txt index 93d4f6ef46..8ab1254fb0 100644 --- a/third_party/ascend/CMakeLists.txt +++ b/third_party/ascend/CMakeLists.txt @@ -137,6 +137,7 @@ add_triton_plugin(TritonAscend ${CMAKE_CURRENT_SOURCE_DIR}/ascend_ir.cc LINK_LIBS AutoBlockifyV1 + TTIRLayoutMerge TritonToLinalg BiShengIRScopeDialect BiShengIRHIVMDialect diff --git a/third_party/ascend/ascend_ir.cc b/third_party/ascend/ascend_ir.cc index 34843f493c..638d1f8475 100644 --- a/third_party/ascend/ascend_ir.cc +++ b/third_party/ascend/ascend_ir.cc @@ -29,7 +29,7 @@ #include "triton/Dialect/Triton/IR/Dialect.h" -#include "AscendModel/RouteModel/SimtSelection.h" +#include "AscendModel/Transforms/SimtSelection.h" #include "ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.h" #include "bishengir/Dialect/Annotation/IR/Annotation.h" #include "bishengir/Dialect/HIVM/IR/HIVM.h" @@ -625,6 +625,26 @@ void init_ascend_ir(py::module &&m) { m.def("remove_attr", [](OpState &op, std::string &name) -> void { op->removeAttr(name); }); + m.def("set_simt_scope_superblock_factor", + [](OpState &root, int64_t factor) -> int64_t { + if (factor != 1 && factor != 2 && factor != 4) + throw std::invalid_argument( + "SIMT scope SuperBlock factor must be 1, 2 or 4"); + int64_t updated = 0; + root->walk([&](Operation *op) { + if (op->getName().getStringRef() != "scope.scope") + return; + auto mode = op->getAttrOfType("vector_mode"); + if (!mode || mode.getValue() != "simt") + return; + OpBuilder builder(op->getContext()); + op->setAttr("ascend.scope_superblock.factor", + builder.getI64IntegerAttr(factor)); + ++updated; + }); + return updated; + }); + m.def("clear_simd_simt_costmodel_attrs", [](OpState &op) { for (llvm::StringRef name : { "ascend.simt_costmodel.effective", diff --git a/third_party/ascend/backend/compiler.py b/third_party/ascend/backend/compiler.py index 4f8f812140..8e7c53bf7d 100644 --- a/third_party/ascend/backend/compiler.py +++ b/third_party/ascend/backend/compiler.py @@ -175,7 +175,8 @@ def _costmodel_profiles_dir() -> Path: return source if source.is_dir() else native_packaged -def _apply_cpp_simd_simt_decision(metadata, effective: str, superblock_factor: int, report: str) -> None: +def _apply_cpp_simd_simt_decision(metadata, effective: str, superblock_factor: int, report: str, + base_num_warps: int = 1) -> None: """Translate the native decision into the backend execution contract.""" metadata["auto_simt_scope_report"] = report metadata["auto_simt_superblock_factor"] = superblock_factor @@ -194,6 +195,16 @@ def _apply_cpp_simd_simt_decision(metadata, effective: str, superblock_factor: i metadata["auto_blockify_v1_runtime_cap"] = False elif effective == "mixed_simd_simt": metadata["auto_simt_requested_kind"] = effective + # Mixed execution still needs V1 at factor one: its persistent + # logical-program loop distributes the launch grid across physical + # cores and is independent of whether F2/F4 SuperBlock aggregation is + # profitable. F2/F4 remain legal candidates, but must not be used as + # the admission condition for V1 itself. + scope_factor = max(1, int(superblock_factor)) + v1_materializable = bool(metadata.get("route_transform_v1_materializable", False)) + metadata["auto_blockify_v1_enabled"] = v1_materializable + metadata["auto_blockify_v1_runtime_cap"] = v1_materializable + metadata["auto_simt_scope_superblock_factor"] = scope_factor elif effective == "all_simt_only": # Auto route selection owns this decision. When the user did not # explicitly disable V1, make the selected pure-SIMT F1/F2/F4 plan @@ -202,7 +213,111 @@ def _apply_cpp_simd_simt_decision(metadata, effective: str, superblock_factor: i metadata["auto_blockify_v1_enabled"] = True -def _run_cpp_simd_simt_costmodel(mod, metadata, opt) -> str: +def _selected_npuir_superblock_factor(metadata, opt) -> int: + """Return the factor owned by the active route. + + CostModel-selected routes publish their factor in metadata. Explicit + manual routes do not run the CostModel and therefore must retain the + launch option. Reading ``auto_simt_scope_superblock_factor`` + unconditionally made an old/default F1 value shadow explicit F2/F4. + """ + effective = metadata.get("auto_simt_effective_kind") + if effective == "mixed_simd_simt" or (effective is None and metadata.get("compile_mode") == "simd_simt"): + # The process-wide factor owns the surrounding SIMD/AIC scheduling + # graph and must stay F1. The selected F2/F4 value is carried on the + # scope itself and restored only in its outlined SIMT module. + selected = 1 + elif effective == "all_simt_only": + selected = metadata.get("auto_simt_superblock_factor", opt.superblock_factor) + else: + selected = opt.superblock_factor + return max(1, int(selected or 1)) + + +def _can_materialize_scope_superblock(metadata, opt, whole_kernel_materializable: bool) -> bool: + """Whether the mixed backend can execute a selected F2/F4 plan. + + The selected Stage is materialized as ``scope.scope``. + NPUIR batches F logical programs around that scope after bufferization, + while the process-wide V1 factor remains one. + """ + del metadata + return bool(whole_kernel_materializable and getattr(opt, "compile_on_910_95", False) + and int(getattr(opt, "num_warps", 0) or 0) > 0) + + +def _publish_route_transform_capability(metadata, opt) -> str: + """Publish the single capability fact consumed by scoring and lowering. + + Layout transforms and the V1 policy have already run when this function is + called. The returned JSON is passed unchanged to the native Route Model so + the report, candidate legality, and backend materialization describe the + same transformed TTIR. + """ + target_supported = bool(getattr(opt, "compile_on_910_95", False)) + num_warps = max(1, int(getattr(opt, "num_warps", 1) or 1)) + v1_enabled = bool(metadata.get("auto_blockify_v1_enabled", False)) + v1_materializable = target_supported and v1_enabled + + disable_reasons = metadata.get("auto_blockify_v1_disable_reasons", []) + if not isinstance(disable_reasons, list): + disable_reasons = [str(disable_reasons)] + if not v1_enabled and not disable_reasons: + disable_reasons = ["not_requested_or_explicitly_disabled"] + if not target_supported: + disable_reasons.append("target_does_not_support_simt_auto_blockify_v1") + + legal_factors = [factor for factor in (1, 2, 4) if num_warps * factor <= 64] + if not legal_factors: + legal_factors = [1] + superblock_factors = legal_factors if v1_materializable else [1] + coalesce_factor = max(1, int(metadata.get("ttir_layout_coalesce_factor", 1) or 1)) + + coalesce_axis = metadata.get("ttir_layout_coalesce_axis", -1) + if not isinstance(coalesce_axis, int): + coalesce_axis = -1 + physical_vector_cores = max(0, int(getattr(opt, "physical_vector_core_count_hint", 0) or 0)) + if not physical_vector_cores: + try: + physical_vector_cores = int(NPUUtils().get_aivector_core_num()) + except Exception: + physical_vector_cores = 0 + source_logical_program_count = max(0, int(getattr(opt, "logical_program_count_hint", 0) or 0)) + transformed_logical_program_count = ((source_logical_program_count + coalesce_factor - 1) // + coalesce_factor if source_logical_program_count else 0) + capability = { + "schema_version": 1, + "layout_merge_applied": bool(metadata.get("ttir_layout_merge_applied", False)), + "row_coalescing_applied": coalesce_factor > 1, + "row_coalescing_factor": coalesce_factor, + "row_coalescing_axis": coalesce_axis, + "auto_blockify_v1_requested": bool(metadata.get("auto_blockify_v1_requested", False)), + "auto_blockify_v1_materializable": v1_materializable, + "auto_blockify_v1_disable_reasons": sorted(set(disable_reasons)), + "whole_kernel_superblock_factors": superblock_factors, + "scope_superblock_factors": superblock_factors, + "source_logical_program_count_hint": source_logical_program_count, + "logical_program_count_hint": transformed_logical_program_count, + "physical_vector_core_count_hint": physical_vector_cores, + } + logical_program_count = capability["logical_program_count_hint"] + if logical_program_count: + capability["superblock_runtime_groups"] = { + str(factor): { + "full_group_count": logical_program_count // factor, + "tail_count": logical_program_count % factor, + } + for factor in (1, 2, 4) + } + capability_json = json.dumps(capability, sort_keys=True, separators=(",", ":")) + metadata["route_transform_capability"] = capability_json + metadata["route_transform_v1_materializable"] = v1_materializable + metadata["route_transform_whole_kernel_factors"] = ",".join(map(str, superblock_factors)) + metadata["route_transform_scope_factors"] = ",".join(map(str, superblock_factors)) + return capability_json + + +def _run_cpp_simd_simt_costmodel(mod, metadata, opt, analysis_ttir_code: str = "") -> str: """Run native selection/materialization; Python only schedules the passes.""" mode = opt.auto_simt_scope_mode if mode == "off" or metadata.get("compile_mode") != "simd_simt": @@ -212,26 +327,37 @@ def _run_cpp_simd_simt_costmodel(mod, metadata, opt) -> str: _costmodel_profiles_dir() / "simd_simt" / "david_v100_simd_simt_v1.json") pm = ir.pass_manager(mod.context) pm.enable_debug() - whole_kernel_superblock_materializable = bool(opt.compile_on_910_95 and opt.enable_auto_blockify is not False) + capability_json = metadata.get("route_transform_capability") + if not capability_json: + capability_json = _publish_route_transform_capability(metadata, opt) + whole_kernel_superblock_materializable = bool(metadata.get("route_transform_v1_materializable", False)) metadata["auto_simt_whole_kernel_superblock_materializable"] = (whole_kernel_superblock_materializable) - ascend.passes.ttir.add_select_simd_simt_costmodel( - pm, - mode, - profile, - str(opt.arch), - int(opt.num_warps), - bool(opt.compile_on_910_95), - whole_kernel_superblock_materializable, - # Whole-kernel AutoBlockify V1 cannot materialize a SuperBlock for a - # local mixed scope. Claiming otherwise lets the Route Model select - # F2/F4 even though the executable still runs the scope as F1. A - # future ScopeSuperBlock pass must flip this only after it can batch - # the SIMD producer, local SIMT scope and SIMD consumer together. - False, - str(opt.auto_simt_scope_dump), - ) - ascend.passes.ttir.add_materialize_simt_scopes(pm) - pm.run(mod, "select_simd_simt_costmodel") + # Mixed F2/F4 requires an outer factor-one V1 loop that NPUIR can refine + # around the selected local scope. + scope_superblock_materializable = _can_materialize_scope_superblock(metadata, opt, + whole_kernel_superblock_materializable) + metadata["auto_simt_scope_superblock_materializable"] = scope_superblock_materializable + with tempfile.TemporaryDirectory() as tmpdir: + analysis_path = "" + if analysis_ttir_code: + analysis_path = os.path.join(tmpdir, "post_auto_blockify_v1.ttir.mlir") + Path(analysis_path).write_text(analysis_ttir_code, encoding="utf-8", newline="\n") + ascend.passes.ttir.add_select_simd_simt_costmodel( + pm, + mode, + profile, + str(opt.arch), + int(opt.num_warps), + bool(opt.compile_on_910_95), + whole_kernel_superblock_materializable, + scope_superblock_materializable, + int(json.loads(capability_json).get("logical_program_count_hint", 0)), + analysis_path, + capability_json, + str(opt.auto_simt_scope_dump), + ) + ascend.passes.ttir.add_materialize_simt_scopes(pm) + pm.run(mod, "select_simd_simt_costmodel") report = ascend.ir.get_string_attr(mod, "ascend.simt_costmodel.report_json") effective = ascend.ir.get_string_attr(mod, "ascend.simt_costmodel.effective") @@ -241,7 +367,7 @@ def _run_cpp_simd_simt_costmodel(mod, metadata, opt) -> str: superblock_factor = 1 if not report or effective not in {"all_simd", "all_simt_only", "mixed_simd_simt", "backend_default"}: raise RuntimeError("invalid native SIMD/SIMT costmodel result") - _apply_cpp_simd_simt_decision(metadata, effective, superblock_factor, report) + _apply_cpp_simd_simt_decision(metadata, effective, superblock_factor, report, int(opt.num_warps)) return effective @@ -277,6 +403,7 @@ def _resolve_auto_blockify_v1_policy(ttir_code: str, metadata, opt) -> bool: env_enabled = _is_auto_map_parallel_blocks_enabled() requested = opt.enable_auto_blockify requested_enabled = env_enabled if requested is None else bool(requested) + metadata["auto_blockify_v1_requested"] = requested_enabled blacklist_reasons = [] has_blacklist_op = metadata.get("has_auto_blockify_blacklist_op") @@ -293,6 +420,14 @@ def _resolve_auto_blockify_v1_policy(ttir_code: str, metadata, opt) -> bool: enabled = requested_enabled and not has_blacklist_op metadata["auto_blockify_v1_enabled"] = enabled + if blacklist_reasons: + metadata["auto_blockify_v1_disable_reasons"] = list(blacklist_reasons) + elif has_blacklist_op: + metadata["auto_blockify_v1_disable_reasons"] = ["explicit_blacklist_override"] + elif not requested_enabled: + metadata["auto_blockify_v1_disable_reasons"] = ["not_requested_or_explicitly_disabled"] + else: + metadata["auto_blockify_v1_disable_reasons"] = [] metadata["auto_blockify_v1_selection_source"] = ("option" if requested is not None else "TRITON_ALL_BLOCKS_PARALLEL") return enabled @@ -325,12 +460,48 @@ def _run_ta_simt_auto_blockify_v1(mod, metadata, opt, *, super_block_factor=None return materialized -def _parse_ttir_text(ttir_code: str, context): +def _parse_ttir_text(ttir_code: str, context=None): """Recreate a TTIR module in ``context`` from an immutable snapshot.""" + if context is None: + context = ir.context() + ir.load_dialects(context) + buffer_ir.load_dialects(context) + ascend_ir.load_dialects(context) + ascend.load_dialects(context) + if distributed is not None: + distributed.ir.load_dialects(context) with tempfile.TemporaryDirectory() as tmpdir: source = Path(tmpdir) / "route_neutral.ttir.mlir" - source.write_text(ttir_code) - return ir.parse_mlir_module(str(source), context) + source.write_text(ttir_code, encoding="utf-8", newline="\n") + module = ir.parse_mlir_module(str(source), context) + module.context = context + return module + + +def _build_costmodel_analysis_ttir(mod, metadata, opt) -> str: + """Build a disposable post-AutoBlockify V1-F1 analysis view. + + The real module stays post-layout and route-neutral. The native Route + Model scores this clone, then materializes its decision on the real + module. Pure-SIMT and mixed lowering apply their executable V1 schedule + only after selection, so all-SIMD never inherits SIMT dispatch IR. + """ + if not metadata.get("auto_blockify_v1_enabled", False): + metadata["auto_simt_costmodel_analysis_ir"] = "post_layout_ttir" + return "" + + analysis_mod = _parse_ttir_text(str(mod), getattr(mod, "context", None)) + analysis_metadata = {} + materialized = _run_ta_simt_auto_blockify_v1( + analysis_mod, + analysis_metadata, + opt, + super_block_factor=1, + ) + metadata["auto_simt_costmodel_analysis_v1_materialized"] = materialized + metadata["auto_simt_costmodel_analysis_ir"] = ("post_auto_blockify_v1_f1_ttir" + if materialized else "post_layout_ttir") + return str(analysis_mod) if materialized else "" def _refine_ta_simt_auto_blockify_v1_superblock(mod, metadata, super_block_factor) -> None: @@ -347,10 +518,19 @@ def _refine_ta_simt_auto_blockify_v1_superblock(mod, metadata, super_block_facto def ttir_to_linalg(mod, metadata, opt, *, named_ops=False): # use triton_adapter to lower Triton-MLIR to linalg + analysis_ttir_code = "" if metadata.get("compile_mode") == "simd_simt" and opt.auto_simt_scope_mode != "off": - _run_ttir_layout_merge(mod, metadata) + if opt.enable_ttir_layout_merge: + _run_ttir_layout_merge(mod, metadata) + else: + metadata["ttir_layout_merge_applied"] = False + metadata["ttir_layout_coalesce_factor"] = 1 + metadata["ttir_layout_coalesce_axis"] = -1 + metadata["ttir_layout_coalesce_grid_ceil_div"] = False _resolve_auto_blockify_v1_policy(str(mod), metadata, opt) - cpp_decision = _run_cpp_simd_simt_costmodel(mod, metadata, opt) + _publish_route_transform_capability(metadata, opt) + analysis_ttir_code = _build_costmodel_analysis_ttir(mod, metadata, opt) + cpp_decision = _run_cpp_simd_simt_costmodel(mod, metadata, opt, analysis_ttir_code) cpp_all_simt = cpp_decision == "all_simt_only" if metadata.get("compile_mode") == "simd_simt" and (cpp_all_simt or ascend.ir.is_whole_body_void_simt_scope(mod)): metadata["scope_pure_simt_auto"] = True @@ -384,15 +564,31 @@ def ttir_to_linalg(mod, metadata, opt, *, named_ops=False): _get_then_remove_rc(mod, "ta.ttir_layout_merge.applied") return str(mod) + # A hand-written mixed kernel already owns its scope boundary, so no Route + # Model pass runs to attach the requested factor. Publish the same local + # scope ABI here; automatic routes attach it during C++ materialization. + if cpp_decision == "backend_default" and metadata.get("compile_mode") == "simd_simt": + manual_scope_factor = max(1, int(opt.superblock_factor or 1)) + manual_scope_count = ascend.ir.set_simt_scope_superblock_factor(mod, manual_scope_factor) + if manual_scope_count: + metadata["auto_simt_scope_superblock_factor"] = manual_scope_factor + metadata["manual_simt_scope_count"] = manual_scope_count + # Get TTIR after C++ has materialized any selected mixed-mode scopes. ttir_code = str(mod) auto_map_parallel_blocks_enabled = _is_auto_map_parallel_blocks_enabled() if metadata.get("auto_simt_effective_kind") == "all_simd": # Preserve the decision-owned disable from # _apply_cpp_simd_simt_decision. Re-running the option policy here - # would re-enable the SIMT-only V1 transform on an all-SIMD binary. + # would re-enable the whole-kernel V1 transform on an all-SIMD binary. metadata["auto_blockify_v1_enabled"] = False metadata["auto_blockify_v1_runtime_cap"] = False + elif metadata.get("auto_simt_effective_kind") == "mixed_simd_simt": + # Preserve the mixed SuperBlock contract established above. NPUIR V1 + # creates the persistent logical-program loop even for F1; the + # selected factor independently controls logical-program aggregation. + metadata["auto_blockify_v1_enabled"] = True + metadata["auto_blockify_v1_runtime_cap"] = True else: _resolve_auto_blockify_v1_policy(ttir_code, metadata, opt) has_auto_blockify_blacklist_op = metadata["has_auto_blockify_blacklist_op"] @@ -1039,18 +1235,8 @@ def linalg_to_bin_enable_npu_compile_910_95(linalg: str, metadata, opt): and metadata.get("parallel_mode") == "mix_simd_simt") if npuir_v1_mixed_enabled: _compile_option_list += ["--enable-auto-blockify-loop"] - # Scope SuperBlock keeps the outer V1 logical-program wrapper and - # runs the selected local SIMT Stage with factor warp groups. The - # surrounding SIMD producer/consumer phases remain inside the V1 - # logical-program body, so values never cross between programs. - if metadata.get("parallel_mode") == "mix_simd_simt": - scope_factor = int(metadata.get( - "auto_simt_superblock_factor", - opt.superblock_factor, - ) or 1) - if scope_factor > 1: - _compile_option_list += [f"--super-block-factor={scope_factor}"] - metadata["scope_superblock_factor"] = scope_factor + selected_factor = _selected_npuir_superblock_factor(metadata, opt) + _compile_option_list += [f"--super-block-factor={selected_factor}"] metadata["auto_blockify_v1_runtime_cap"] = True elif not metadata.get("ta_auto_blockify_v1_materialized", False): metadata["auto_blockify_v1_runtime_cap"] = False @@ -1333,7 +1519,11 @@ def linalg_to_bin_enable_npu_compile_A2_A3(linalg: str, metadata, opt): and not metadata.get("ta_auto_blockify_v1_materialized", False) and metadata.get("parallel_mode") == "mix_simd_simt") if npuir_v1_mixed_enabled: - _compile_option_list += ["--enable-auto-blockify-loop"] + selected_factor = _selected_npuir_superblock_factor(metadata, opt) + _compile_option_list += [ + "--enable-auto-blockify-loop", + f"--super-block-factor={selected_factor}", + ] metadata["auto_blockify_v1_runtime_cap"] = True elif not metadata.get("ta_auto_blockify_v1_materialized", False): metadata["auto_blockify_v1_runtime_cap"] = False @@ -1544,6 +1734,14 @@ class NPUOptions: # take effect on the reorder instruction pattern for SIMT. The pattern is disabled by default. enable_simt_reorder_instruction: bool = False enable_costmodel_backend: bool = False + # Apply implicit-permute and row/tile/strided coalescing before Route Model + # analysis. Exposed primarily for controlled compiler ablation; normal + # compilation should retain the default. + enable_ttir_layout_merge: bool = True + # Optional runtime launch fact used by SuperBlock costing. Zero means the + # frontend cannot provide a stable logical-program count for this compile. + logical_program_count_hint: int = 0 + physical_vector_core_count_hint: int = 0 # Native AscendModel selection: Python only schedules the C++ passes. auto_simt_scope_mode: str = "" auto_simt_scope_dump: str = "" @@ -1606,7 +1804,13 @@ def ttir_to_npubin(mod, metadata, opt): # Pure-SIMT must observe the same post-layout TTIR as the Route # Model. In particular this restores the independent-row merge # before AutoBlockify V1 creates its persistent scheduling loop. - _run_ttir_layout_merge(mod, metadata) + if opt.enable_ttir_layout_merge: + _run_ttir_layout_merge(mod, metadata) + else: + metadata["ttir_layout_merge_applied"] = False + metadata["ttir_layout_coalesce_factor"] = 1 + metadata["ttir_layout_coalesce_axis"] = -1 + metadata["ttir_layout_coalesce_grid_ceil_div"] = False auto_blockify_v1_enabled = _resolve_auto_blockify_v1_policy(str(mod), metadata, opt) if opt.enable_ta_auto_blockify_v1 and auto_blockify_v1_enabled: materialized = _run_ta_simt_auto_blockify_v1(mod, metadata, opt) @@ -1663,9 +1867,8 @@ def ttir_to_npubin(mod, metadata, opt): # port's behavior can be validated without a hidden fallback. if (not opt.enable_ta_auto_blockify_v1 and metadata.get("auto_blockify_v1_enabled", False)): _compile_option_list += ["--enable-auto-blockify-loop"] - selected_factor = int(metadata.get("auto_simt_superblock_factor", opt.superblock_factor) or 1) - if selected_factor > 0: - _compile_option_list += [f"--super-block-factor={selected_factor}"] + selected_factor = _selected_npuir_superblock_factor(metadata, opt) + _compile_option_list += [f"--super-block-factor={selected_factor}"] npu_compiler_path, env = _get_npucompiler_path() cmd_list = ([npu_compiler_path, src_path] + _compile_option_list + ["-o", bin_file]) diff --git a/third_party/ascend/backend/utils.py b/third_party/ascend/backend/utils.py index 0786aacdfc..da74a162fa 100644 --- a/third_party/ascend/backend/utils.py +++ b/third_party/ascend/backend/utils.py @@ -50,6 +50,9 @@ def is_compile_on_910_95(): return _is_compile_on_910_95 +# Cache/eviction modifiers are deliberately not blacklist reasons. V1 only +# changes logical-program scheduling and both the TTIR row lift and V1 cloning +# preserve those attributes on the rebuilt load/store operations. AUTO_BLOCKIFY_BLACKLIST_RULES = ( (re.compile(r"\btt\.atomic_(?:rmw|cas)\b"), "atomic operations"), (re.compile(r"\btt\.elementwise_inline_asm\b"), "inline elementwise assembly"), @@ -57,10 +60,6 @@ def is_compile_on_910_95(): re.compile(r"\btt\.load\b[^\n]*\bisVolatile\s*=\s*true\b"), "loads with volatile", ), - ( - re.compile(r"\btt\.(?:load|store)\b[^\n]*\bcacheModifier\s*="), - "loads or stores with cache modifiers", - ), ) backend_policy = None diff --git a/third_party/ascend/costmodel/include/AscendModel/Analysis/HIVMAnalysis.h b/third_party/ascend/costmodel/include/AscendModel/Analysis/HIVMAnalysis.h index 73393486f2..7642e01a11 100644 --- a/third_party/ascend/costmodel/include/AscendModel/Analysis/HIVMAnalysis.h +++ b/third_party/ascend/costmodel/include/AscendModel/Analysis/HIVMAnalysis.h @@ -26,8 +26,6 @@ namespace mlir { namespace ascend { -enum class HIVMSchedulerMode { Static, DES }; - enum class HIVMPipe { Unknown, Vector, @@ -71,7 +69,6 @@ struct HIVMOp { struct HIVMAnalysisReport { std::string sourcePath; std::string sourceMode; - HIVMSchedulerMode schedulerMode = HIVMSchedulerMode::Static; int64_t oneIterationCycles = 0; int64_t weightedCycles = 0; int64_t totalBusyCycles = 0; @@ -89,16 +86,12 @@ struct HIVMAnalysisReport { void print(llvm::raw_ostream &os, const HardwareConfig &config) const; void emitPerfettoTrace(llvm::raw_ostream &os, const HardwareConfig &config) const; - void emitDESGraph(llvm::raw_ostream &os, const HardwareConfig &config) const; - void emitFeedbackJSON(llvm::raw_ostream &os, - const HardwareConfig &config) const; }; class HIVMAnalyzer { public: HIVMAnalyzer(const HardwareConfig &config, - llvm::StringRef argBindings = llvm::StringRef(), - HIVMSchedulerMode schedulerMode = HIVMSchedulerMode::Static); + llvm::StringRef argBindings = llvm::StringRef()); bool analyzeModule(mlir::ModuleOp module, HIVMAnalysisReport &report, std::string &error) const; @@ -107,12 +100,10 @@ class HIVMAnalyzer { std::string &error) const; static llvm::StringRef stringifyPipe(HIVMPipe pipe); - static llvm::StringRef stringifySchedulerMode(HIVMSchedulerMode mode); private: const HardwareConfig &config; std::string argBindingsStr; - HIVMSchedulerMode schedulerMode; }; } // namespace ascend diff --git a/third_party/ascend/costmodel/include/AscendModel/RouteModel/SimtAnchorAnalysis.h b/third_party/ascend/costmodel/include/AscendModel/Analysis/SimtAnchorAnalysis.h similarity index 64% rename from third_party/ascend/costmodel/include/AscendModel/RouteModel/SimtAnchorAnalysis.h rename to third_party/ascend/costmodel/include/AscendModel/Analysis/SimtAnchorAnalysis.h index fa2e0dc42d..ade6465105 100644 --- a/third_party/ascend/costmodel/include/AscendModel/RouteModel/SimtAnchorAnalysis.h +++ b/third_party/ascend/costmodel/include/AscendModel/Analysis/SimtAnchorAnalysis.h @@ -16,7 +16,6 @@ #include #include #include -#include #include namespace mlir { @@ -36,52 +35,10 @@ enum class SimtAnchorKind { llvm::StringRef stringifySimtAnchorKind(SimtAnchorKind kind); -/// Whether one physical route is available for an anchor. BackendConditional -/// may be reported and scored offline, but it is not production-selectable. -enum class CandidateLoweringStatus { - Unsupported, - Native, - BackendConditional, - AliasesMixed, -}; - -llvm::StringRef -stringifyCandidateLoweringStatus(CandidateLoweringStatus status); - struct CandidateLowerability { - CandidateLoweringStatus allSimd = CandidateLoweringStatus::Native; - CandidateLoweringStatus allSimtOnly = CandidateLoweringStatus::Native; - CandidateLoweringStatus mixed = CandidateLoweringStatus::Native; - std::vector allSimdReasons; - std::vector allSimtOnlyReasons; - std::vector mixedReasons; -}; - -struct TensorAtomicFacts { - int64_t updateElements = 0; - int64_t addressRank = 0; - std::string valueType = "unknown"; - std::string offsetType = "unknown"; - std::string operation = "unknown"; - bool hasMask = false; - std::optional staticMaskActiveFraction; - bool resultUsed = false; - bool addressIsLaneVarying = false; - bool addressDependsOnLoadedIndex = false; - std::string contention = "unknown"; -}; - -struct HistogramFacts { - int64_t inputElements = 0; - int64_t numBins = 0; - std::string inputType = "unknown"; - std::string resultType = "unknown"; -}; - -struct PlainCumsumFacts { - int64_t axisExtent = 0; - std::string elementType = "unknown"; - bool reverse = false; + bool allSimd = true; + bool allSimtOnly = true; + bool mixed = true; }; /// Structural facts for a blockwise triangular recurrence such as solve_tril. @@ -98,10 +55,6 @@ struct TriangularSolveFacts { bool requiresCubeTailPartition = false; }; -using SimtAnchorFacts = - std::variant; - struct SimtAnchorDescriptor { Operation *operation = nullptr; /// Exact top-level TTIR operations that will be moved into one local SIMT @@ -116,7 +69,7 @@ struct SimtAnchorDescriptor { /// initial SIMD loads to reproduce the hand-written solve_tril scope. Operation *scopeInsertionPoint = nullptr; SimtAnchorKind kind = SimtAnchorKind::LoadedIndexDependentMemory; - SimtAnchorFacts facts; + std::optional triangularSolve; CandidateLowerability lowerability; /// True only when the current target/materializer contract can turn this /// descriptor into a local SIMT scope. @@ -131,20 +84,26 @@ struct SimtAnchorPlan { CandidateLowerability kernelLowerability; llvm::SmallVector materializableRoots() const; - int64_t materializableCount() const; }; +/// Merge the materializable anchors owned by one LogicalStage into the exact +/// compound scope that will be scored and materialized. When several anchor +/// operations share a block, every operation in the lexical interval between +/// the first and last anchor is included so operands are never captured from +/// after the new scope. A failure means the Stage cannot be represented by +/// one local SIMT scope. +std::optional +mergeSimtStageAnchors(const SimtAnchorPlan &plan, + llvm::ArrayRef anchorIndices); + /// Materialize exactly the local SIMT regions described by `plan`. /// /// This is deliberately a transform over the immutable analysis result: it /// does not rediscover anchors, recompute features, or read per-operation /// selection attributes. The caller owns the final effective route decision. LogicalResult materializeSimtAnchorPlan(ModuleOp module, - const SimtAnchorPlan &plan); - -/// Classify a TTIR operation by SIMT mechanism, independently of whether the -/// selected target can currently materialize it. -std::optional classifyMixedSimtAnchor(Operation *op); + const SimtAnchorPlan &plan, + int64_t superblockFactor = 1); /// True when a load/store pointer has an SSA backward slice that reaches a /// loaded/gathered index. This is a real data-dependence test and must not be @@ -154,15 +113,6 @@ bool isLoadedIndexDependentMemoryOp(Operation *op); /// Build the non-overlapping shared plan in pre-order. SimtAnchorPlan buildMixedSimtAnchorPlan(ModuleOp module, bool compileOn91095); -/// Return true when `op` is a TTIR operation that the current local-scope -/// materializer can move into a SIMT scope on 910_95-class targets. -bool isMixedSimtAnchor(Operation *op, bool compileOn91095); - -/// Collect non-overlapping anchors in pre-order. Nested operations are not -/// returned when an enclosing operation already forms one materialized scope. -llvm::SmallVector collectMixedSimtAnchors(ModuleOp module, - bool compileOn91095); - } // namespace ascend } // namespace mlir diff --git a/third_party/ascend/costmodel/include/AscendModel/RouteModel/StagePartitioner.h b/third_party/ascend/costmodel/include/AscendModel/Analysis/StagePartitioner.h similarity index 80% rename from third_party/ascend/costmodel/include/AscendModel/RouteModel/StagePartitioner.h rename to third_party/ascend/costmodel/include/AscendModel/Analysis/StagePartitioner.h index e44a21484d..fb1d31b7ca 100644 --- a/third_party/ascend/costmodel/include/AscendModel/RouteModel/StagePartitioner.h +++ b/third_party/ascend/costmodel/include/AscendModel/Analysis/StagePartitioner.h @@ -1,7 +1,7 @@ //===- StagePartitioner.h - Build semantic Phase/Stage IR ----*- C++ -*-===// -#ifndef ASCENDMODEL_ROUTEMODEL_STAGEPARTITIONER_H -#define ASCENDMODEL_ROUTEMODEL_STAGEPARTITIONER_H +#ifndef ASCENDMODEL_ANALYSIS_STAGEPARTITIONER_H +#define ASCENDMODEL_ANALYSIS_STAGEPARTITIONER_H #include "AscendModel/RouteModel/SimdSimtCostModel.h" #include "AscendModel/RouteModel/StageCostModels.h" @@ -68,10 +68,6 @@ class ProgramStructureAnalysis { /// target implementation for production Stage ownership. class PhaseBoundaryAnalysis { public: - llvm::Expected> - analyze(const SimdSimtFeatureSummary &features, - const StagePartitionerOptions &options) const; - llvm::Expected> analyze(ModuleOp module, const SimtAnchorPlan &anchorPlan, const SimdSimtFeatureSummary &features, @@ -111,8 +107,6 @@ class StageKindClassifier { class StageWorkloadAnalysis { public: llvm::Error analyze(StagePartition &partition) const; - llvm::Error verify(const StagePartition &partition, - const StageWorkload &kernelWorkload) const; }; /// Derives legal SIMD/SIMT implementations from structural Stage facts. @@ -125,8 +119,7 @@ class StageModeLegalityAnalysis { class StagePartitionVerifier { public: - llvm::Error verify(const StagePartition &partition, - const StageWorkload &kernelWorkload) const; + llvm::Error verify(const StagePartition &partition) const; }; /// Partitions post-layout/post-AutoBlockify-V1 TTIR facts into serial Phases @@ -134,25 +127,12 @@ class StagePartitionVerifier { /// but it is never internally mixed. class StagePartitioner { public: - llvm::Expected> - partition(const SimdSimtFeatureSummary &features, - const StagePartitionerOptions &options) const; - - llvm::Expected> - partition(const SimdSimtFeatureSummary &features, - const StagePartitionerOptions &options, - const SimtAnchorPlan &anchorPlan) const; - llvm::Expected> partition(ModuleOp module, const SimtAnchorPlan &anchorPlan, const SimdSimtFeatureSummary &features, const StagePartitionerOptions &options) const; }; -/// Builds the mode-independent kernel workload used only for ownership -/// conservation checks. It contains logical elements/bytes, never cycles. -StageWorkload buildKernelStageWorkload(const SimdSimtFeatureSummary &features); - } // namespace mlir::ascend -#endif // ASCENDMODEL_ROUTEMODEL_STAGEPARTITIONER_H +#endif // ASCENDMODEL_ANALYSIS_STAGEPARTITIONER_H diff --git a/third_party/ascend/costmodel/include/AscendModel/RouteModel/SimdSimtCostModel.h b/third_party/ascend/costmodel/include/AscendModel/RouteModel/SimdSimtCostModel.h index daefe7ac01..6991cabf28 100644 --- a/third_party/ascend/costmodel/include/AscendModel/RouteModel/SimdSimtCostModel.h +++ b/third_party/ascend/costmodel/include/AscendModel/RouteModel/SimdSimtCostModel.h @@ -8,15 +8,13 @@ #ifndef ASCENDMODEL_ROUTEMODEL_SIMDSIMTCOSTMODEL_H #define ASCENDMODEL_ROUTEMODEL_SIMDSIMTCOSTMODEL_H -#include "AscendModel/RouteModel/SimtAnchorAnalysis.h" +#include "AscendModel/Analysis/SimtAnchorAnalysis.h" #include "AscendModel/RouteModel/StageRouteCostModel.h" #include "mlir/IR/BuiltinOps.h" -#include "llvm/ADT/StringMap.h" #include "llvm/Support/Error.h" #include "llvm/Support/JSON.h" #include "llvm/Support/raw_ostream.h" -#include #include #include #include @@ -36,63 +34,10 @@ llvm::StringRef stringifySimdSimtCandidate(SimdSimtCandidateKind candidate); /// current materializer can wrap in local SIMT scopes. It intentionally /// contains no Operation pointers, so reports remain stable and serializable. struct SimtAnchorFeatureSummary { - /// All mechanism anchors recognized before target/lowering checks. - int64_t recognizedCount = 0; - /// Materializable anchors used by the mixed candidate. `count` is retained - /// as the serialized compatibility name for this value. int64_t count = 0; - int64_t coveredOperationCount = 0; - int64_t loadOps = 0; - int64_t storeOps = 0; - int64_t reduceOps = 0; - int64_t scanOps = 0; - int64_t gatherOps = 0; - int64_t dotOps = 0; - int64_t atomicOps = 0; - int64_t histogramOps = 0; - - int64_t maxTensorNumel = 1; - int64_t maxElementBits = 0; - int64_t maskRankSum = 0; - int64_t uniqueMaskValues = 0; - int64_t uniqueMaskRankSum = 0; - int64_t predicateElements = 0; - /// Loop-weighted mask lanes on predicate-producing/consuming IR edges. - int64_t predicateLaneEvaluations = 0; - int64_t pointerTensorOps = 0; - int64_t loadedIndexDependentMemoryOps = 0; - int64_t laneDependentPointerOps = 0; - int64_t maxReduceAxisExtent = 1; - int64_t weightedReduceAxisElements = 0; - /// Loop-weighted input lanes times log2(reduction/scan axis extent). - int64_t shuffleLaneSteps = 0; - int64_t staticLoopCount = 0; - int64_t staticLoopTripCountSum = 0; - int64_t modeledDynamicLoopCount = 0; - int64_t modeledDynamicLoopTripCountSum = 0; int64_t conditionalBranchCount = 0; int64_t divergentBranchCount = 0; double activeLaneRatio = 1.0; - - bool hasControlFlow = false; - - llvm::StringMap weightedOps; - llvm::StringMap opElements; - double loadBytes = 0.0; - double storeBytes = 0.0; - int64_t loadWarpInstructions = 0; - int64_t storeWarpInstructions = 0; - int64_t dotFlops = 0; - - int64_t capturedTensorCount = 0; - int64_t escapingTensorCount = 0; - double capturedTensorBytes = 0.0; - double escapingTensorBytes = 0.0; - - std::vector mechanismKinds; - std::vector tensorAtomics; - std::vector histograms; - std::vector plainCumsums; std::vector triangularSolves; CandidateLowerability kernelLowerability; @@ -106,81 +51,10 @@ struct SimdSimtFeatureSummary { int64_t loadOps = 0; int64_t storeOps = 0; int64_t reduceOps = 0; - int64_t scanOps = 0; - int64_t gatherOps = 0; int64_t dotOps = 0; - int64_t atomicOps = 0; - int64_t histogramOps = 0; - int64_t broadcastOps = 0; - int64_t expandDimsOps = 0; - int64_t splatOps = 0; - int64_t addPtrOps = 0; - - int64_t arithOps = 0; - int64_t mathOps = 0; - int64_t addOps = 0; - int64_t subOps = 0; - int64_t mulOps = 0; - int64_t divOps = 0; - int64_t maxOps = 0; - int64_t absOps = 0; - int64_t expOps = 0; - int64_t logOps = 0; - int64_t cmpOps = 0; - int64_t selectOps = 0; - int64_t castOps = 0; - int64_t clampOps = 0; - int64_t scalarOps = 0; - - int64_t maxTensorRank = 0; - int64_t maxTensorNumel = 1; - int64_t maxElementBits = 0; - int64_t maskTensorOps = 0; - int64_t maskRankSum = 0; - int64_t uniqueMaskValues = 0; - int64_t uniqueMaskRankSum = 0; - int64_t predicateElements = 0; - /// Loop-weighted mask lanes on predicate-producing/consuming IR edges. - int64_t predicateLaneEvaluations = 0; - int64_t maskBroadcastOps = 0; - int64_t pointerTensorOps = 0; - int64_t pointerUnstructuredDims = 0; - /// Real SSA provenance count. Unlike laneDependentPointerOps, this field - /// requires the address backward slice to reach a loaded/gathered index. int64_t loadedIndexDependentMemoryOps = 0; - /// Legacy rank-based proxy retained for report compatibility. - int64_t laneDependentPointerOps = 0; - int64_t rowLocalReduceOps = 0; - int64_t maxReduceAxisExtent = 1; - int64_t weightedReduceAxisElements = 0; - /// Loop-weighted input lanes times log2(reduction/scan axis extent). - int64_t shuffleLaneSteps = 0; - int64_t scalarLoadOps = 0; - int64_t scalarStoreOps = 0; - int64_t vectorPtrSplatOps = 0; - int64_t vectorReduceToScalarOps = 0; - bool rank1IndirectVectorReduce = false; - - llvm::StringMap weightedOps; - llvm::StringMap opElements; - double loadBytes = 0.0; - double storeBytes = 0.0; - int64_t loadWarpInstructions = 0; - int64_t storeWarpInstructions = 0; int64_t dotFlops = 0; - int64_t dotOutputElements = 0; - std::vector> dotMNK; - int64_t staticLoopCount = 0; - int64_t staticLoopTripCountSum = 0; int64_t staticLoopTripCountMax = 1; - int64_t modeledDynamicLoopCount = 0; - int64_t modeledDynamicLoopTripCountSum = 0; - /// scf.for iter_args whose values feed data computation in a later - /// iteration. These dependencies serialize a Stage roofline. - int64_t loopCarriedDataDependencyCount = 0; - /// Loop-carried pointer/address induction tracked separately because - /// downstream address lowering can remove it. - int64_t pointerInductionDependencyCount = 0; int64_t conditionalBranchCount = 0; int64_t divergentBranchCount = 0; double activeLaneRatio = 1.0; @@ -188,122 +62,23 @@ struct SimdSimtFeatureSummary { /// Scheduling/layout facts read from the transformed TTIR consumed by this /// model. These make it explicit that layout merging and AutoBlockify V1 /// ran before feature extraction rather than being guessed from source TTIR. - bool ttirLayoutMergeApplied = false; - int64_t coalesceFactor = 1; - int64_t coalesceAxis = -1; bool autoBlockifyV1Applied = false; int64_t autoBlockifyV1LoopCount = 0; - int64_t autoBlockifyV1ScheduleOpCount = 0; - bool autoBlockifyV1HasDynamicTripCount = false; - - bool hasDot = false; - bool hasGather = false; - bool hasAtomic = false; - bool hasHistogram = false; - bool hasScan = false; bool hasExplicitScope = false; - bool hasControlFlow = false; - bool hasDynamicShape = false; - bool hasUnknownTripCount = false; - - /// Pattern observations are diagnostic inputs to the model. They do not - /// create a mandatory online route: all three candidates stay selectable. - std::vector observedMixedKinds; SimtAnchorFeatureSummary simtAnchors; llvm::json::Object toJSON() const; }; -/// Applicability is strictly a hardware/lowering statement. It says whether -/// transformed TTIR contains a recognized SIMT mechanism and whether the -/// current target can materialize at least one corresponding anchor. -struct SimtApplicabilityResult { - bool mechanismDetected = false; - bool targetSupported = false; - bool materializable = false; - int64_t recognizedAnchorCount = 0; - int64_t materializableAnchorCount = 0; - std::vector mechanisms; - std::vector reasons; - - llvm::json::Object toJSON() const; -}; - struct SimdSimtCandidateScores { double allSimd = 0.0; double allSimtOnly = 0.0; double mixedSimdSimt = 0.0; - double get(SimdSimtCandidateKind candidate) const; llvm::json::Object toJSON() const; }; -/// Detailed values retained so the JSON report can explain every major term -/// in the versioned analytical and structural formulas. -struct SimdSimtCostBreakdown { - llvm::StringMap simdOpSystemCycles; - llvm::StringMap simtOpSystemCycles; - llvm::StringMap structuralComponents; - - double simdComputeCycles = 0.0; - double simtComputeCycles = 0.0; - double simdDotCycles = 0.0; - double simtDotCycles = 0.0; - double simdLoadCycles = 0.0; - double simdStoreCycles = 0.0; - double simdMemoryCycles = 0.0; - double simtLoadCycles = 0.0; - double simtStoreCycles = 0.0; - double simtMemoryCycles = 0.0; - double simtShuffleInstructions = 0.0; - double simtShuffleCycles = 0.0; - double simtPredicateInstructions = 0.0; - double simtPredicateCycles = 0.0; - double simdSetupCycles = 0.0; - double simtSetupCycles = 0.0; - double simdIssuePayloadCycles = 0.0; - double simtIssuePayloadCycles = 0.0; - double simdAnalyticalCycles = 0.0; - double simtAnalyticalCycles = 0.0; - double programIssueScale = 1.0; - - double mixedSimdRegularComputeCycles = 0.0; - double mixedSimdRegularDotCycles = 0.0; - double mixedSimdRegularMemoryCycles = 0.0; - double mixedSimdRegularPayloadCycles = 0.0; - double mixedSimtAnchorComputeCycles = 0.0; - double mixedSimtAnchorDotCycles = 0.0; - double mixedSimtAnchorMemoryCycles = 0.0; - double mixedSimtAnchorShuffleCycles = 0.0; - double mixedSimtAnchorPredicateCycles = 0.0; - double mixedSimtAnchorPayloadCycles = 0.0; - double mixedSimtAnchorCalibratedPayloadCycles = 0.0; - int64_t cubeTailDotOps = 0; - int64_t cubeTailDotFlops = 0; - double mixedBoundaryCycles = 0.0; - double mixedRemainingStructuralPenaltyRatio = 0.0; - - double irregularDensity = 0.0; - double tinyDotUnderfill = 0.0; - double structuralPenaltyRatio = 0.0; - /// SIMD-only residual cost for structures omitted by its compute/memory - /// roofline. It must not depend on the cost of another route candidate. - double simdStructuralPenaltyCycles = 0.0; - - double mixedSimdFraction = 0.0; - /// Conservative low-confidence setup fallback for a mixed plan. The - /// current source is a standalone empty-VF harness, not a measured - /// directional SIMD/SIMT transition. - double mixedSetupFallbackCycles = 0.0; - double standaloneSimtSetupCycles = 0.0; - double setupProxyDeltaCycles = 0.0; - int64_t mixedSetupFallbackNumWarps = 0; - std::string mixedCostSource; - - llvm::json::Object toJSON(const SimdSimtFeatureSummary &features) const; -}; - struct SimdSimtCostModelOptions { /// Empty selects TRITON_ASCEND_SIMD_SIMT_PROFILE, then the source-tree /// profile compiled into AscendModelRouteModel. @@ -321,6 +96,11 @@ struct SimdSimtCostModelOptions { /// True when backend integration can apply AutoBlockify V1 to a pure-SIMT /// kernel. This is deliberately independent of local-scope batching. bool wholeKernelSuperblockMaterializable = false; + /// Optional runtime launch count. Zero means unknown; a positive value + /// prevents the solver from pricing factors that cannot form one full + /// logical-program group. + int64_t logicalProgramCountHint = 0; + int64_t physicalVectorCoreCountHint = 0; }; struct SimdSimtCostReport { @@ -335,22 +115,14 @@ struct SimdSimtCostReport { std::string microbenchmarkProfileTarget; std::string microbenchmarkProfileContentSha256; std::string scoreUnit; - std::string scoreScope = "per_program_ranking_proxy"; - SimdSimtCandidateScores candidateCosts; - SimdSimtCandidateScores candidateRatiosToBest; bool allSimdCandidateLegal = true; bool allSimtOnlyCandidateLegal = true; bool mixedCandidateLegal = false; SimdSimtCandidateKind decision = SimdSimtCandidateKind::AllSIMD; - double bestScore = 0.0; - bool targetCompatible = true; std::vector unsupported; - SimtApplicabilityResult applicability; - SimdSimtFeatureSummary features; - SimdSimtCostBreakdown breakdown; StageCostModelSummary stageModel; bool includeFeaturesInJSON = true; @@ -373,14 +145,6 @@ llvm::Expected analyzeSimdSimtFeatures(mlir::ModuleOp module, const SimtAnchorPlan &anchorPlan); -/// Run the versioned profile formula on an already materialized feature -/// summary. Every structurally lowerable candidate is scored; policy-domain, -/// confidence, Event-validation, and gain-margin admission are deliberately -/// absent from the online model. -llvm::Expected -estimateSimdSimtCandidates(const SimdSimtFeatureSummary &features, - const SimdSimtCostModelOptions &options = {}); - /// Analyze a ModuleOp and score all three candidates in one call. llvm::Expected analyzeSimdSimtCandidates(mlir::ModuleOp module, diff --git a/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageCostModels.h b/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageCostModels.h index e497e36e75..c1f3aaa097 100644 --- a/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageCostModels.h +++ b/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageCostModels.h @@ -17,8 +17,6 @@ #include "llvm/Support/Error.h" #include -#include -#include #include #include @@ -48,7 +46,6 @@ enum class StageCostModelKind { }; llvm::StringRef stringifyStageCostModel(StageCostModelKind kind); -std::optional parseStageCostModel(llvm::StringRef name); struct StageControlFlowRates { double loopBackedgeCycles = 0.0; @@ -61,7 +58,6 @@ struct StageControlFlowRates { struct LogicalStage { std::string id; - std::string description; StageCostModelKind costModelKind = StageCostModelKind::ScalarIssue; StageScheduleKind scheduleKind = StageScheduleKind::StraightLine; int64_t iterationCount = 1; @@ -99,15 +95,11 @@ struct LogicalStage { struct LogicalPhase { std::string id; - std::string description; std::vector stages; }; struct StagePartition { std::string domain; - /// "operation_graph" for exact post-transform TTIR ownership, otherwise - /// "feature_summary_fallback" for the temporary aggregate implementation. - std::string boundarySource = "feature_summary_fallback"; bool operationOwnershipComplete = false; int64_t modeledOperationCount = 0; std::vector phases; @@ -170,67 +162,10 @@ struct HardwareProfile { bool isValid() const; }; -class ProfileProvider { -public: - explicit ProfileProvider(HardwareProfile profile); - - llvm::Expected - getSnapshot(llvm::StringRef target, llvm::StringRef profileVersion) const; - -private: - HardwareProfile profile; -}; - -struct StageCostModelContext { - const LogicalStage &stage; - const HardwareProfile &profile; -}; - -class StageCostModel { -public: - virtual ~StageCostModel() = default; - virtual StageMode getMode() const = 0; - virtual llvm::StringRef getName() const = 0; - virtual bool supports(StageCostModelKind kind) const = 0; - virtual double estimate(const StageCostModelContext &context, - const StageImplementation &implementation, - const StageResourceCycles &resources) const = 0; -}; - -class SIMDStageCostModel : public StageCostModel { -public: - StageMode getMode() const final { return StageMode::SIMD; } -}; - -class SIMTStageCostModel : public StageCostModel { -public: - StageMode getMode() const final { return StageMode::SIMT; } -}; - -class StageCostModelRegistry { -public: - static const StageCostModelRegistry &get(); - - llvm::Expected lookup(StageMode mode, - StageCostModelKind kind) const; - llvm::Error verifyComplete() const; - -private: - StageCostModelRegistry(); - std::vector> models; -}; - class StageCostEvaluator { public: - explicit StageCostEvaluator( - const StageCostModelRegistry ®istry = StageCostModelRegistry::get()) - : registry(registry) {} - llvm::Expected evaluate(const StagePartition &partition, const HardwareProfile &profile) const; - -private: - const StageCostModelRegistry ®istry; }; } // namespace mlir::ascend diff --git a/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageRouteCostModel.h b/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageRouteCostModel.h index 36fa5f1e90..92a3dab1f2 100644 --- a/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageRouteCostModel.h +++ b/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageRouteCostModel.h @@ -30,13 +30,16 @@ enum class StageScheduleKind { }; llvm::StringRef stringifyStageMode(StageMode mode); -llvm::StringRef stringifyStageKernelRoute(StageKernelRouteKind kind); -llvm::StringRef stringifyStageSchedule(StageScheduleKind kind); struct StageImplementation { StageMode mode = StageMode::SIMD; - /// SIMD always uses factor=1. SIMT may use F1/F2/F4 when legal. + /// SIMD always uses factor=1. For a whole-kernel SIMT implementation this + /// is the AutoBlockify V1 factor. A local SIMT implementation identifies a + /// mixed-kernel candidate whose selected Stage is materialized as a scope. + /// The current backend still applies factor>1 through the surrounding V1 + /// kernel schedule; it is not an independently widened scope VF. int64_t superblockFactor = 1; + bool localScope = false; bool isValid() const; llvm::json::Object toJSON() const; @@ -64,7 +67,6 @@ struct StageModelFeatures { /// a single recurrence use one group. int64_t parallelRecurrenceGroupCount = 1; double activeLaneRatio = 1.0; - std::string source; bool isValid() const; bool permitsSimdRoofline() const; @@ -92,8 +94,8 @@ struct StageWorkload { }; /// Resource costs for one iteration after raw Stage workload has been mapped -/// through the selected immutable hardware profile. setup/epilogue are paid -/// once; all other fields are per iteration. +/// through the selected immutable hardware profile. Setup is paid once; all +/// other fields are per iteration. struct StageResourceCycles { double setup = 0.0; double scalar = 0.0; @@ -103,7 +105,6 @@ struct StageResourceCycles { double predicate = 0.0; double shuffle = 0.0; double dot = 0.0; - double control = 0.0; double loopControl = 0.0; double branchControl = 0.0; double divergence = 0.0; @@ -111,7 +112,6 @@ struct StageResourceCycles { double spill = 0.0; double issue = 0.0; double criticalPath = 0.0; - double epilogue = 0.0; bool isFiniteAndNonNegative() const; llvm::json::Object toJSON() const; @@ -121,9 +121,6 @@ struct StageImplementationCost { StageImplementation implementation; double totalCycles = 0.0; StageResourceCycles resources; - std::string modelName; - std::string profileVersion; - std::string source; bool isValid() const; llvm::json::Object toJSON() const; @@ -131,7 +128,6 @@ struct StageImplementationCost { struct LogicalStageCost { std::string id; - std::string description; std::string model; StageScheduleKind schedule = StageScheduleKind::StraightLine; int64_t iterationCount = 1; @@ -153,6 +149,9 @@ struct LogicalStageCost { int64_t scopeOutputTensorBytes = 0; std::vector simtAnchorIndices; bool localSimtMaterializable = false; + /// Factors legal for a whole-kernel pure-SIMT schedule. + std::vector legalSimtFactors; + /// Factors legal when this Stage alone is materialized as a local scope. std::vector localSimtFactors; std::vector implementations; @@ -161,7 +160,6 @@ struct LogicalStageCost { struct LogicalPhaseCost { std::string id; - std::string description; std::vector stages; llvm::json::Object toJSON() const; @@ -169,14 +167,13 @@ struct LogicalPhaseCost { struct StageCostTable { std::string domain; - std::string boundarySource; bool operationOwnershipComplete = false; int64_t modeledOperationCount = 0; std::string profileVersion; + int64_t logicalProgramCountHint = 0; + int64_t physicalCoreCountHint = 0; std::vector phases; std::vector stages; - - llvm::json::Object toJSON() const; }; struct StageTransitionCost { @@ -190,7 +187,6 @@ struct StageTransitionCost { double simtUbLoadBytesPerThreadPerCycle = 1.0; double simtUbStoreBytesPerThreadPerCycle = 1.0; int64_t simtWarpSize = 1; - std::string source; bool isValid() const; double get(StageMode from, StageMode to) const; @@ -205,8 +201,9 @@ struct StageRoutePlan { std::vector logicalStageCycles; std::vector logicalPhaseCycles; int64_t routeSuperblockFactor = 1; + int64_t runtimePhysicalProgramCount = 0; + int64_t runtimeWaveCount = 1; double totalCycles = 0.0; - std::string source; llvm::json::Object toJSON() const; }; @@ -214,7 +211,6 @@ struct StageRoutePlan { struct StageCostModelSummary { bool applied = false; std::string domain; - std::string boundarySource; bool operationOwnershipComplete = false; int64_t modeledOperationCount = 0; std::string profileVersion; diff --git a/third_party/ascend/costmodel/include/AscendModel/Transforms/Passes.td b/third_party/ascend/costmodel/include/AscendModel/Transforms/Passes.td index 180f048c45..cd90bda15d 100644 --- a/third_party/ascend/costmodel/include/AscendModel/Transforms/Passes.td +++ b/third_party/ascend/costmodel/include/AscendModel/Transforms/Passes.td @@ -154,9 +154,6 @@ def HIVMAnalysisPass : Pass<"analyze-hivm", "ModuleOp"> { }]; let options = [ - Option<"schedulerMode", "scheduler", "std::string", - /*default=*/"\"static\"", - "Scheduler backend: static or des">, Option<"argBindingsStr", "arg-bindings", "std::string", /*default=*/"\"\"", "Bindings for dynamic HIVM values (e.g., 'arg10=128,arg11=64')">, @@ -225,6 +222,15 @@ def SelectSimdSimtCostModelPass Option<"scopeSuperblockMaterializable", "scope-superblock-materializable", "bool", /*default=*/"false", "Whether backend integration can materialize mixed F2/F4 with AutoBlockify V1">, + Option<"logicalProgramCountHint", "logical-program-count-hint", "int64_t", + /*default=*/"0", + "Optional runtime logical-program count used to cost full SuperBlock groups and tails">, + Option<"analysisModulePath", "analysis-module-path", "std::string", + /*default=*/"\"\"", + "Optional post-layout/post-AutoBlockify TTIR used only for feature analysis and scoring">, + Option<"routeTransformCapabilityJSON", "route-transform-capability-json", "std::string", + /*default=*/"\"{}\"", + "Resolved layout/V1 capability shared by scoring, materialization, and reporting">, Option<"reportFile", "report-file", "std::string", /*default=*/"\"\"", "Optional JSONL path for the C++ selection report"> diff --git a/third_party/ascend/costmodel/include/AscendModel/RouteModel/SimtSelection.h b/third_party/ascend/costmodel/include/AscendModel/Transforms/SimtSelection.h similarity index 100% rename from third_party/ascend/costmodel/include/AscendModel/RouteModel/SimtSelection.h rename to third_party/ascend/costmodel/include/AscendModel/Transforms/SimtSelection.h diff --git a/third_party/ascend/costmodel/lib/AscendModel/Analysis/HIVMAnalysis.cpp b/third_party/ascend/costmodel/lib/AscendModel/Analysis/HIVMAnalysis.cpp index 6f1d20b2fe..47e566de9f 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/Analysis/HIVMAnalysis.cpp +++ b/third_party/ascend/costmodel/lib/AscendModel/Analysis/HIVMAnalysis.cpp @@ -154,17 +154,6 @@ struct EventKey { } }; -struct EventInstanceKey { - EventKey key; - int64_t generation = 0; - - bool operator<(const EventInstanceKey &other) const { - return std::tie(key.sender, key.receiver, key.eventId, generation) < - std::tie(other.key.sender, other.key.receiver, other.key.eventId, - other.generation); - } -}; - struct LoopFrame { int braceDepth = 0; int64_t tripCount = 1; @@ -2305,405 +2294,8 @@ static void finalizeScheduledReport(HIVMAnalysisReport &report, report.weightedCycles = report.oneIterationCycles; } -struct CompletionEvent { - int64_t time = 0; - size_t opId = 0; - - bool operator>(const CompletionEvent &other) const { - return std::tie(time, opId) > std::tie(other.time, other.opId); - } -}; - -struct BufferSlotState { - int64_t writableAt = 0; - int64_t readableAt = 0; - int64_t version = 0; -}; - -struct BufferRootState { - std::vector slots; - int64_t latestReadableAt = 0; - int64_t latestVersion = 0; - std::map versionReadableAt; - std::map versionToSlot; -}; - -static llvm::StringRef getSyncBlockSourceCore(const HIVMOp &op) { - bool isCubeCore = op.coreType == "CUBE" || op.coreType == "AIC"; - if (op.opName == "sync_block_set") - return isCubeCore ? "AIC" : "AIV"; - if (op.opName == "sync_block_wait") - return isCubeCore ? "AIV" : "AIC"; - return ""; -} - -static void normalizeSyncBlockGenerations(HIVMAnalysisReport &report) { - std::map, int64_t> setGeneration; - std::map, int64_t> waitGeneration; - for (HIVMOp &op : report.operations) { - if ((op.opName != "sync_block_set" && op.opName != "sync_block_wait") || - op.eventId.empty()) - continue; - llvm::StringRef sourceCore = getSyncBlockSourceCore(op); - if (sourceCore.empty()) - continue; - auto key = std::make_pair(op.eventId, sourceCore.str()); - if (op.opName == "sync_block_set") - op.eventGeneration = ++setGeneration[key]; - else - op.eventGeneration = ++waitGeneration[key]; - } -} - -/// After generation normalization, wire explicit dependency edges from each -/// sync_block_set to its matching sync_block_wait so the DES respects -/// cross-core ordering. Without this, the wait may be scheduled before the -/// set completes (they live in different func::FuncOps with independent state). -static void wireCrossCoreSyncDependencies(HIVMAnalysisReport &report) { - // Key: (eventId, sourceCore, generation) → set-op id - using SyncKey = std::tuple; - std::map setOpById; - for (HIVMOp &op : report.operations) { - if (op.opName != "sync_block_set" || op.eventId.empty()) - continue; - llvm::StringRef sourceCore = getSyncBlockSourceCore(op); - if (sourceCore.empty()) - continue; - SyncKey key{op.eventId, sourceCore.str(), op.eventGeneration}; - setOpById[key] = op.id; - } - for (HIVMOp &op : report.operations) { - if (op.opName != "sync_block_wait" || op.eventId.empty()) - continue; - llvm::StringRef sourceCore = getSyncBlockSourceCore(op); - if (sourceCore.empty()) - continue; - // sync_block_wait's sourceCore returns the core that *set* the flag - // (the opposite core), which matches the set-op's sourceCore. - SyncKey key{op.eventId, sourceCore.str(), op.eventGeneration}; - auto it = setOpById.find(key); - if (it != setOpById.end()) - op.dependsOn.push_back(it->second); - } -} - -static void finalizeDiscreteEventReport(HIVMAnalysisReport &report, - const HardwareConfig &config) { - normalizeSyncBlockGenerations(report); - wireCrossCoreSyncDependencies(report); - - const size_t numOps = report.operations.size(); - if (numOps == 0) { - report.weightedCycles = 0; - return; - } - - std::vector remainingDeps(numOps, 0); - std::vector readyAt(numOps, 0); - std::vector> successors(numOps); - std::vector queued(numOps, false); - std::vector started(numOps, false); - std::vector completed(numOps, false); - std::deque readyOps; - std::priority_queue, - std::greater> - completions; - std::map pipeAvailableAt; - std::map flagEventVisibleAt; - std::map blockSyncVisibleAt; - std::map bufferStates; - std::map>> - writeSlotAssignments; - size_t completedCount = 0; - - for (const HIVMOp &op : report.operations) { - for (const std::string &root : op.writeBuffers) { - auto &state = bufferStates[root]; - if (state.slots.empty()) { - int64_t count = std::max(1, op.multiBufferSlots); - for (int64_t i = 0; i < count; ++i) - state.slots.push_back(BufferSlotState{}); - } - state.versionReadableAt.emplace(0, 0); - } - for (const std::string &root : op.readBuffers) { - auto [it, inserted] = bufferStates.try_emplace(root, BufferRootState{}); - it->second.versionReadableAt.emplace(0, 0); - } - } - - for (size_t opId = 0; opId < numOps; ++opId) { - HIVMOp &op = report.operations[opId]; - remainingDeps[opId] = op.dependsOn.size(); - for (size_t depId : op.dependsOn) { - if (depId < numOps) - successors[depId].push_back(opId); - } - if (remainingDeps[opId] == 0) { - readyOps.push_back(opId); - queued[opId] = true; - } - } - - auto completeOp = [&](size_t opId, int64_t time) { - if (completed[opId]) - return; - HIVMOp &op = report.operations[opId]; - op.endCycle = time; - completed[opId] = true; - ++completedCount; - report.oneIterationCycles = - std::max(report.oneIterationCycles, op.endCycle); - report.totalBusyCycles += op.duration; - report.opCount++; - if (op.pipe == HIVMPipe::Scalar && op.opName != "set_flag" && - op.opName != "wait_flag" && op.opName != "sync_block_set" && - op.opName != "sync_block_wait" && op.opName != "sync_block" && - op.opName != "pipe_barrier" && op.opName != "get_block_idx" && - op.opName != "get_block_num" && op.opName != "get_sub_block_idx" && - op.opName != "get_sub_block_num" && op.opName != "set_mask_norm" && - op.opName != "pointer_cast" && op.opName != "convert_layout") - report.unknownOpCount++; - if (op.isSyncOp) { - report.syncCycles += op.duration; - report.syncOpCount++; - if ((op.opName == "set_flag" || op.opName == "sync_block_set") && - !op.eventId.empty()) { - EventInstanceKey key{{op.senderPipe, op.receiverPipe, op.eventId}, - op.eventGeneration}; - if (op.opName == "sync_block_set") - blockSyncVisibleAt[key] = time; - else - flagEventVisibleAt[key] = time; - } - } - auto slotIt = writeSlotAssignments.find(opId); - if (slotIt != writeSlotAssignments.end()) { - for (const auto &[root, slotIndex] : slotIt->second) { - auto rootIt = bufferStates.find(root); - if (rootIt == bufferStates.end() || - slotIndex >= rootIt->second.slots.size()) - continue; - BufferRootState &state = rootIt->second; - BufferSlotState &slot = state.slots[slotIndex]; - auto rootVersionIt = llvm::find(op.writeBuffers, root); - if (rootVersionIt == op.writeBuffers.end()) - continue; - size_t bufferIdx = - std::distance(op.writeBuffers.begin(), rootVersionIt); - if (bufferIdx >= op.writeBufferVersions.size()) - continue; - int64_t version = op.writeBufferVersions[bufferIdx]; - slot.readableAt = time; - slot.version = version; - state.latestVersion = std::max(state.latestVersion, version); - state.latestReadableAt = std::max(state.latestReadableAt, time); - state.versionReadableAt[version] = time; - state.versionToSlot[version] = slotIndex; - } - } - if (op.isBarrier) { - report.barrierCycles += op.duration; - report.barrierCount++; - } - if (op.pipe != HIVMPipe::All && op.pipe != HIVMPipe::Unknown) { - report.pipeBusyCycles[op.pipe] += op.duration; - report.weightedPipeCycles[op.pipe] += op.duration * op.loopMultiplier; - } - for (size_t succId : successors[opId]) { - readyAt[succId] = std::max(readyAt[succId], time); - if (remainingDeps[succId] > 0) - --remainingDeps[succId]; - if (remainingDeps[succId] == 0 && !queued[succId]) { - readyOps.push_back(succId); - queued[succId] = true; - } - } - }; - - auto computeStartTime = [&](const HIVMOp &op) -> int64_t { - int64_t start = readyAt[op.id]; - for (size_t idx = 0; idx < op.readBuffers.size(); ++idx) { - const std::string &root = op.readBuffers[idx]; - auto it = bufferStates.find(root); - if (it == bufferStates.end()) - continue; - int64_t version = - idx < op.readBufferVersions.size() ? op.readBufferVersions[idx] : 0; - auto readableIt = it->second.versionReadableAt.find(version); - if (readableIt != it->second.versionReadableAt.end()) - start = std::max(start, readableIt->second); - else if (it->second.latestVersion >= version) - start = std::max(start, it->second.latestReadableAt); - } - for (const std::string &root : op.writeBuffers) { - auto it = bufferStates.find(root); - if (it != bufferStates.end() && !it->second.slots.empty()) { - int64_t slotReady = std::numeric_limits::max(); - for (const BufferSlotState &slot : it->second.slots) - slotReady = std::min(slotReady, slot.writableAt); - start = std::max(start, slotReady); - } - } - if ((op.opName == "wait_flag" || op.opName == "sync_block_wait") && - !op.eventId.empty()) { - EventInstanceKey key{{op.senderPipe, op.receiverPipe, op.eventId}, - op.eventGeneration}; - auto &visibleAt = op.opName == "sync_block_wait" ? blockSyncVisibleAt - : flagEventVisibleAt; - auto it = visibleAt.find(key); - if (it != visibleAt.end()) - start = std::max(start, it->second); - } - if (op.pipe == HIVMPipe::Unknown) - return start; - if (op.isBarrier && op.pipe == HIVMPipe::All) { - if (op.coreType.empty()) { - for (const auto &entry : pipeAvailableAt) - start = std::max(start, entry.second); - } else { - for (const auto &entry : pipeAvailableAt) { - if (pipeBelongsToCore(entry.first, op.coreType)) - start = std::max(start, entry.second); - } - } - return start; - } - return std::max(start, pipeAvailableAt[op.pipe]); - }; - - auto startOp = [&](size_t opId, int64_t startTime) { - HIVMOp &op = report.operations[opId]; - started[opId] = true; - op.startCycle = startTime; - const int64_t endTime = startTime + op.duration; - for (size_t idx = 0; idx < op.readBuffers.size(); ++idx) { - const std::string &root = op.readBuffers[idx]; - auto it = bufferStates.find(root); - if (it == bufferStates.end()) - continue; - int64_t version = - idx < op.readBufferVersions.size() ? op.readBufferVersions[idx] : 0; - if (version <= 0) - continue; - auto slotIt = it->second.versionToSlot.find(version); - if (slotIt == it->second.versionToSlot.end()) - continue; - size_t slotIndex = slotIt->second; - if (slotIndex >= it->second.slots.size()) - continue; - it->second.slots[slotIndex].writableAt = - std::max(it->second.slots[slotIndex].writableAt, endTime); - } - for (const std::string &root : op.writeBuffers) { - auto it = bufferStates.find(root); - if (it == bufferStates.end() || it->second.slots.empty()) - continue; - BufferRootState &state = it->second; - size_t bestSlot = 0; - int64_t bestTime = state.slots.front().writableAt; - for (size_t i = 1; i < state.slots.size(); ++i) { - if (state.slots[i].writableAt < bestTime) { - bestTime = state.slots[i].writableAt; - bestSlot = i; - } - } - state.slots[bestSlot].writableAt = endTime; - writeSlotAssignments[opId].push_back({root, bestSlot}); - } - if (op.pipe != HIVMPipe::Unknown) { - if (op.isBarrier && op.pipe == HIVMPipe::All) { - auto barrierPipes = getCoreBarrierPipes(op.coreType); - if (barrierPipes.empty()) { - for (auto &entry : pipeAvailableAt) - entry.second = endTime; - } else { - for (HIVMPipe barrierPipe : barrierPipes) - pipeAvailableAt[barrierPipe] = endTime; - } - } else { - pipeAvailableAt[op.pipe] = endTime; - } - } - if (op.duration == 0) - completeOp(opId, endTime); - else - completions.push({endTime, opId}); - }; - - int64_t currentTime = 0; - while (completedCount < numOps) { - bool startedAny = false; - size_t readyCount = readyOps.size(); - for (size_t i = 0; i < readyCount; ++i) { - size_t opId = readyOps.front(); - readyOps.pop_front(); - HIVMOp &op = report.operations[opId]; - if (started[opId] || completed[opId]) - continue; - int64_t startTime = computeStartTime(op); - if (startTime <= currentTime) { - startOp(opId, currentTime); - startedAny = true; - } else { - readyOps.push_back(opId); - } - } - - while (!completions.empty() && completions.top().time <= currentTime) { - size_t opId = completions.top().opId; - completions.pop(); - if (!completed[opId]) { - completeOp(opId, currentTime); - } - } - - if (startedAny) - continue; - - int64_t nextTime = std::numeric_limits::max(); - if (!completions.empty()) - nextTime = std::min(nextTime, completions.top().time); - for (size_t opId : readyOps) - nextTime = std::min(nextTime, computeStartTime(report.operations[opId])); - - if (nextTime == std::numeric_limits::max()) - break; - currentTime = std::max(currentTime, nextTime); - - while (!completions.empty() && completions.top().time <= currentTime) { - size_t opId = completions.top().opId; - completions.pop(); - if (!completed[opId]) { - completeOp(opId, currentTime); - } - } - } - - int64_t globalBarrierWeightedCycles = 0; - for (const HIVMOp &op : report.operations) { - if (op.isBarrier && op.pipe == HIVMPipe::All) - globalBarrierWeightedCycles += op.duration * op.loopMultiplier; - } - for (const auto &entry : report.weightedPipeCycles) - report.weightedCycles = std::max(report.weightedCycles, entry.second); - report.weightedCycles += globalBarrierWeightedCycles; - if (report.weightedCycles == 0) - report.weightedCycles = report.oneIterationCycles; -} - } // namespace -llvm::StringRef HIVMAnalyzer::stringifySchedulerMode(HIVMSchedulerMode mode) { - switch (mode) { - case HIVMSchedulerMode::Static: - return "static"; - case HIVMSchedulerMode::DES: - return "des"; - } - return "static"; -} - llvm::StringRef HIVMAnalyzer::stringifyPipe(HIVMPipe pipe) { switch (pipe) { case HIVMPipe::Vector: @@ -2731,10 +2323,8 @@ llvm::StringRef HIVMAnalyzer::stringifyPipe(HIVMPipe pipe) { } HIVMAnalyzer::HIVMAnalyzer(const HardwareConfig &config, - llvm::StringRef argBindings, - HIVMSchedulerMode schedulerMode) - : config(config), argBindingsStr(argBindings.str()), - schedulerMode(schedulerMode) {} + llvm::StringRef argBindings) + : config(config), argBindingsStr(argBindings.str()) {} bool HIVMAnalyzer::analyzeFile(llvm::StringRef path, HIVMAnalysisReport &report, std::string &error) const { @@ -2749,7 +2339,6 @@ bool HIVMAnalyzer::analyzeFile(llvm::StringRef path, HIVMAnalysisReport &report, report = HIVMAnalysisReport(); report.sourcePath = path.str(); report.sourceMode = "direct-hivm"; - report.schedulerMode = schedulerMode; { mlir::DialectRegistry registry; @@ -2804,16 +2393,10 @@ bool HIVMAnalyzer::analyzeModule(mlir::ModuleOp module, report = HIVMAnalysisReport(); report.sourcePath = ""; report.sourceMode = "mlir-pass"; - report.schedulerMode = schedulerMode; - AnalysisState state; state.argBindings = parseArgBindings(argBindingsStr); - analyzeParsedRegion(module.getBodyRegion(), 1, state, report, config, - schedulerMode == HIVMSchedulerMode::DES); - if (schedulerMode == HIVMSchedulerMode::DES) - finalizeDiscreteEventReport(report, config); - else - finalizeScheduledReport(report, config); + analyzeParsedRegion(module.getBodyRegion(), 1, state, report, config, false); + finalizeScheduledReport(report, config); return true; } @@ -2822,8 +2405,6 @@ void HIVMAnalysisReport::print(llvm::raw_ostream &os, os << "=== HIVM Analysis ===\n"; os << "Source mode: " << sourceMode << "\n"; os << "Source: " << sourcePath << "\n"; - os << "Scheduler: " << HIVMAnalyzer::stringifySchedulerMode(schedulerMode) - << "\n"; os << "Hardware: " << config.getName() << " @ " << llvm::format("%.2f", config.getClockFrequencyGHz()) << " GHz\n\n"; @@ -2884,58 +2465,6 @@ void HIVMAnalysisReport::print(llvm::raw_ostream &os, } } -void HIVMAnalysisReport::emitFeedbackJSON(llvm::raw_ostream &os, - const HardwareConfig &config) const { - const double fallbackRatio = - opCount > 0 ? static_cast(unknownOpCount) / opCount : 1.0; - const bool scheduleComplete = - opCount > 0 && operations.size() == opCount && unknownOpCount == 0; - - llvm::json::Object pipes; - for (const auto &entry : pipeBusyCycles) - pipes[HIVMAnalyzer::stringifyPipe(entry.first).str()] = entry.second; - - llvm::json::Object weightedPipes; - for (const auto &entry : weightedPipeCycles) - weightedPipes[HIVMAnalyzer::stringifyPipe(entry.first).str()] = - entry.second; - - llvm::json::Object quality; - quality["parsed_operation_count"] = static_cast(operations.size()); - quality["scheduled_operation_count"] = static_cast(opCount); - quality["fallback_operation_count"] = static_cast(unknownOpCount); - quality["fallback_ratio"] = fallbackRatio; - quality["schedule_complete"] = scheduleComplete; - - llvm::json::Object summary; - summary["critical_path_cycles"] = oneIterationCycles; - summary["critical_path_us"] = config.cyclesToMicroseconds(oneIterationCycles); - summary["weighted_cycles"] = weightedCycles; - summary["sync_cycles"] = syncCycles; - summary["barrier_cycles"] = barrierCycles; - summary["operation_count"] = static_cast(opCount); - summary["sync_operation_count"] = static_cast(syncOpCount); - summary["barrier_count"] = static_cast(barrierCount); - summary["max_loop_multiplier"] = maxLoopMultiplier; - - llvm::json::Object root; - root["schema_version"] = 1; - root["kind"] = "hivm_des_feedback_source"; - root["source_mode"] = sourceMode; - root["source"] = sourcePath; - root["scheduler"] = HIVMAnalyzer::stringifySchedulerMode(schedulerMode).str(); - root["target"] = config.getName(); - root["scope"] = "single_program"; - root["cycle_domain"] = "device_clock_cycles"; - root["clock_frequency_ghz"] = config.getClockFrequencyGHz(); - root["excludes"] = llvm::json::Array({"host_launch", "grid_wave_count"}); - root["quality"] = std::move(quality); - root["summary"] = std::move(summary); - root["pipe_busy_cycles"] = std::move(pipes); - root["weighted_pipe_cycles"] = std::move(weightedPipes); - os << llvm::formatv("{0:2}", llvm::json::Value(std::move(root))); -} - void HIVMAnalysisReport::emitPerfettoTrace(llvm::raw_ostream &os, const HardwareConfig &config) const { // Assign each pipe a unique tid. Pipes are grouped into AIC (Cube core) @@ -3164,78 +2693,3 @@ void HIVMAnalysisReport::emitPerfettoTrace(llvm::raw_ostream &os, os << "\n ],\n \"displayTimeUnit\": \"us\"\n}\n"; } - -void HIVMAnalysisReport::emitDESGraph(llvm::raw_ostream &os, - const HardwareConfig &config) const { - auto joinStrVec = [](const std::vector &v) { - std::string s; - llvm::raw_string_ostream ss(s); - ss << "["; - for (size_t i = 0; i < v.size(); ++i) { - if (i) - ss << ","; - ss << "\"" << v[i] << "\""; - } - ss << "]"; - ss.flush(); - return s; - }; - auto joinIntVec = [](const std::vector &v) { - std::string s; - llvm::raw_string_ostream ss(s); - ss << "["; - for (size_t i = 0; i < v.size(); ++i) { - if (i) - ss << ","; - ss << v[i]; - } - ss << "]"; - ss.flush(); - return s; - }; - auto joinSizeVec = [](const std::vector &v) { - std::string s; - llvm::raw_string_ostream ss(s); - ss << "["; - for (size_t i = 0; i < v.size(); ++i) { - if (i) - ss << ","; - ss << v[i]; - } - ss << "]"; - ss.flush(); - return s; - }; - - os << "{\n"; - os << " \"clock_ghz\": " - << llvm::format("%.3f", config.getClockFrequencyGHz()) << ",\n"; - os << " \"operations\": [\n"; - for (size_t i = 0; i < operations.size(); ++i) { - const HIVMOp &op = operations[i]; - if (i) - os << ",\n"; - os << " {" - << "\"id\":" << op.id << ",\"name\":\"" << op.opName << "\"" - << ",\"pipe\":\"" << HIVMAnalyzer::stringifyPipe(op.pipe) << "\"" - << ",\"duration\":" << op.duration << ",\"line\":" << op.lineNumber - << ",\"depends_on\":" << joinSizeVec(op.dependsOn) - << ",\"is_sync\":" << (op.isSyncOp ? "true" : "false") - << ",\"is_barrier\":" << (op.isBarrier ? "true" : "false") - << ",\"event_id\":\"" << op.eventId << "\"" - << ",\"event_generation\":" << op.eventGeneration - << ",\"sender_pipe\":\"" << HIVMAnalyzer::stringifyPipe(op.senderPipe) - << "\"" - << ",\"receiver_pipe\":\"" - << HIVMAnalyzer::stringifyPipe(op.receiverPipe) << "\"" - << ",\"core_type\":\"" << op.coreType << "\"" - << ",\"bytes\":" << op.bytes << ",\"elements\":" << op.elements - << ",\"loop_multiplier\":" << op.loopMultiplier - << ",\"multi_buffer_slots\":" << op.multiBufferSlots - << ",\"read_buffers\":" << joinStrVec(op.readBuffers) - << ",\"write_buffers\":" << joinStrVec(op.writeBuffers) - << ",\"read_versions\":" << joinIntVec(op.readBufferVersions) - << ",\"write_versions\":" << joinIntVec(op.writeBufferVersions) << "}"; - } - os << "\n ]\n}\n"; -} diff --git a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/SimtAnchorAnalysis.cpp b/third_party/ascend/costmodel/lib/AscendModel/Analysis/SimtAnchorAnalysis.cpp similarity index 68% rename from third_party/ascend/costmodel/lib/AscendModel/RouteModel/SimtAnchorAnalysis.cpp rename to third_party/ascend/costmodel/lib/AscendModel/Analysis/SimtAnchorAnalysis.cpp index 526ab893fa..13de0f1bbc 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/SimtAnchorAnalysis.cpp +++ b/third_party/ascend/costmodel/lib/AscendModel/Analysis/SimtAnchorAnalysis.cpp @@ -1,6 +1,6 @@ //===- SimtAnchorAnalysis.cpp - Materializable SIMT anchors --------------===// -#include "AscendModel/RouteModel/SimtAnchorAnalysis.h" +#include "AscendModel/Analysis/SimtAnchorAnalysis.h" #include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/BuiltinTypes.h" @@ -42,7 +42,12 @@ static std::string getScalarTypeName(Type type) { return "unknown"; } -static std::optional +struct PlainCumsumLegality { + int64_t axisExtent; + std::string elementType; +}; + +static std::optional analyzePlainOneDimensionalCumsum(Operation *op) { if (!op || op->getName().getStringRef() != "tt.scan" || op->getNumOperands() != 1 || op->getNumResults() != 1) @@ -80,12 +85,8 @@ analyzePlainOneDimensionalCumsum(Operation *op) { terminator->getName().getStringRef() != "tt.scan.return") return std::nullopt; - PlainCumsumFacts facts; - facts.axisExtent = sourceType.getShape()[axisValue]; - facts.elementType = getScalarTypeName(sourceType.getElementType()); - if (auto reverse = op->getAttrOfType("reverse")) - facts.reverse = reverse.getValue(); - return facts; + return PlainCumsumLegality{sourceType.getShape()[axisValue], + getScalarTypeName(sourceType.getElementType())}; } static bool hasTensorPointerOperand(Operation *op) { @@ -148,31 +149,6 @@ static Value findPointerOffset(Value pointer) { return {}; } -static std::optional getStaticMaskActiveFraction(Value mask) { - if (!mask) - return std::nullopt; - Operation *producer = mask.getDefiningOp(); - if (!producer) - return std::nullopt; - if (producer->getName().getStringRef() == "arith.constant") { - Attribute value = producer->getAttr("value"); - if (auto dense = dyn_cast_or_null(value)) { - if (!dense.getElementType().isInteger(1) || dense.getNumElements() == 0) - return std::nullopt; - int64_t active = 0; - for (bool item : dense.getValues()) - active += item; - return static_cast(active) / dense.getNumElements(); - } - if (auto integer = dyn_cast_or_null(value)) - return integer.getValue().isZero() ? 0.0 : 1.0; - } - if (producer->getName().getStringRef() == "tt.splat" && - producer->getNumOperands() == 1) - return getStaticMaskActiveFraction(producer->getOperand(0)); - return std::nullopt; -} - static std::string getAtomicOperation(Operation *op) { if (op->getName().getStringRef() == "tt.atomic_cas") return "cas"; @@ -488,14 +464,6 @@ collectTriangularSolveScopeOperations(Operation *anchor, return result; } -static CandidateLowerability -simpleMixedLowerability(llvm::StringRef allSimtReason) { - CandidateLowerability result; - result.allSimtOnly = CandidateLoweringStatus::BackendConditional; - result.allSimtOnlyReasons.push_back(allSimtReason.str()); - return result; -} - static std::optional analyzeAnchor(Operation *op, bool compileOn91095) { if (!op) @@ -506,151 +474,77 @@ static std::optional analyzeAnchor(Operation *op, if (name == "tt.gather") { descriptor.kind = SimtAnchorKind::DirectGather; - descriptor.lowerability = simpleMixedLowerability( - "whole_module_pure_simt_requires_backend_check"); } else if (name == "tt.histogram") { descriptor.kind = SimtAnchorKind::Histogram; - HistogramFacts facts; auto input = op->getNumOperands() > 0 ? dyn_cast(op->getOperand(0).getType()) : RankedTensorType(); auto result = op->getNumResults() > 0 ? dyn_cast(op->getResult(0).getType()) : RankedTensorType(); - facts.inputElements = getStaticNumElements(input); - facts.numBins = result && result.hasStaticShape() && result.getRank() == 1 - ? result.getShape()[0] - : 0; - facts.inputType = input ? getScalarTypeName(input) : "unknown"; - facts.resultType = result ? getScalarTypeName(result) : "unknown"; - descriptor.facts = facts; - descriptor.lowerability.allSimd = CandidateLoweringStatus::Unsupported; - descriptor.lowerability.allSimdReasons.push_back( - "ascend950_plain_histogram_aliases_backend_mixed_template"); - descriptor.lowerability.allSimtOnly = CandidateLoweringStatus::Unsupported; - descriptor.lowerability.allSimtOnlyReasons.push_back( - "tt.histogram_not_legalized_by_pure_simt_pipeline"); + const int64_t inputElements = getStaticNumElements(input); + const int64_t numBins = + result && result.hasStaticShape() && result.getRank() == 1 + ? result.getShape()[0] + : 0; + const std::string inputType = input ? getScalarTypeName(input) : "unknown"; + const std::string resultType = + result ? getScalarTypeName(result) : "unknown"; + descriptor.lowerability.allSimd = false; + descriptor.lowerability.allSimtOnly = false; bool supported = input && result && input.hasStaticShape() && result.hasStaticShape() && input.getRank() == 1 && result.getRank() == 1 && - (facts.inputType == "i8" || facts.inputType == "i16" || - facts.inputType == "i32" || facts.inputType == "i64") && - facts.resultType == "i32" && facts.inputElements > 0 && - facts.numBins > 0; - if (!supported) { - descriptor.lowerability.mixed = CandidateLoweringStatus::Unsupported; - descriptor.lowerability.mixedReasons.push_back( - "histogram_requires_static_rank1_integer_input_and_rank1_i32_bins"); - } + (inputType == "i8" || inputType == "i16" || + inputType == "i32" || inputType == "i64") && + resultType == "i32" && inputElements > 0 && numBins > 0; + if (!supported) + descriptor.lowerability.mixed = false; } else if (name == "tt.scan") { auto facts = analyzePlainOneDimensionalCumsum(op); if (!facts) return std::nullopt; descriptor.kind = SimtAnchorKind::PlainOneDimensionalCumsum; - descriptor.facts = *facts; - descriptor.lowerability.allSimd = CandidateLoweringStatus::AliasesMixed; - descriptor.lowerability.allSimdReasons.push_back( - "ascend950_plain_1d_cumsum_symbol_is_simt_template"); - descriptor.lowerability.allSimtOnly = - CandidateLoweringStatus::BackendConditional; - descriptor.lowerability.allSimtOnlyReasons.push_back( - "pure_simt_cumsum_requires_whole_module_backend_check"); - if (facts->axisExtent <= 0 || !isSupportedCumsumType(facts->elementType)) { - descriptor.lowerability.mixed = CandidateLoweringStatus::Unsupported; - descriptor.lowerability.mixedReasons.push_back( - "plain_1d_cumsum_dtype_or_axis_extent_not_supported"); - } else if (facts->axisExtent <= 64) { - descriptor.lowerability.mixedReasons.push_back( - "template_uses_small_register_path_without_simt_async_invoke"); - } + descriptor.lowerability.allSimd = false; + if (facts->axisExtent <= 0 || !isSupportedCumsumType(facts->elementType)) + descriptor.lowerability.mixed = false; } else if (name == "tt.atomic_rmw" || name == "tt.atomic_cas") { descriptor.kind = SimtAnchorKind::TensorAtomic; - TensorAtomicFacts facts; auto result = op->getNumResults() > 0 ? dyn_cast(op->getResult(0).getType()) : RankedTensorType(); - facts.updateElements = getStaticNumElements(result); - facts.valueType = result ? getScalarTypeName(result) : "unknown"; - facts.operation = getAtomicOperation(op); + const int64_t updateElements = getStaticNumElements(result); + const std::string valueType = + result ? getScalarTypeName(result) : "unknown"; + const std::string operation = getAtomicOperation(op); + std::string offsetType = "unknown"; if (op->getNumOperands() > 0) { - if (auto pointer = - dyn_cast(op->getOperand(0).getType())) { - facts.addressRank = pointer.getRank(); - facts.addressIsLaneVarying = getStaticNumElements(pointer) > 1; - } Value offset = findPointerOffset(op->getOperand(0)); - if (offset) { - facts.offsetType = getScalarTypeName(offset.getType()); - if (auto offsetTensor = dyn_cast(offset.getType())) - facts.addressIsLaneVarying |= getStaticNumElements(offsetTensor) > 1; - } - facts.addressDependsOnLoadedIndex = pointerDependsOnLoadedIndex(op); + if (offset) + offsetType = getScalarTypeName(offset.getType()); } - facts.hasMask = name == "tt.atomic_rmw" && op->getNumOperands() >= 3; - if (facts.hasMask) - facts.staticMaskActiveFraction = - getStaticMaskActiveFraction(op->getOperand(2)); - facts.resultUsed = op->getNumResults() > 0 && !op->getResult(0).use_empty(); - facts.contention = "unknown"; - descriptor.facts = facts; - descriptor.lowerability = - simpleMixedLowerability("pure_simt_atomic_requires_backend_check"); - - bool supported = result && result.hasStaticShape() && - facts.updateElements > 0 && - isSupportedAtomicType(facts.valueType, facts.operation) && - (facts.offsetType == "i32" || facts.offsetType == "i64"); - if (facts.resultUsed && - (facts.valueType == "f16" || facts.valueType == "bf16")) { + bool supported = result && result.hasStaticShape() && updateElements > 0 && + isSupportedAtomicType(valueType, operation) && + (offsetType == "i32" || offsetType == "i64"); + const bool resultUsed = + op->getNumResults() > 0 && !op->getResult(0).use_empty(); + if (resultUsed && (valueType == "f16" || valueType == "bf16")) supported = false; - descriptor.lowerability.mixedReasons.push_back( - "f16_bf16_atomic_old_value_semantics_require_validation"); - } - if (!supported) { - descriptor.lowerability.mixed = CandidateLoweringStatus::Unsupported; - descriptor.lowerability.mixedReasons.push_back( - "atomic_shape_dtype_operation_or_offset_not_supported"); - } + if (!supported) + descriptor.lowerability.mixed = false; } else if (isTriangularSolveLoop(op)) { descriptor.kind = SimtAnchorKind::TriangularSolveLoop; - descriptor.facts = analyzeTriangularSolveFacts(op); - descriptor.lowerability = simpleMixedLowerability( - "pure_simt_triangular_solve_requires_cube_tail_partition"); - descriptor.lowerability.mixedReasons.push_back( - "manual_scope_shape_vector16_masked_reduce_loop"); + descriptor.triangularSolve = analyzeTriangularSolveFacts(op); } else if (isLoadedIndexDependentMemoryOp(op)) { descriptor.kind = SimtAnchorKind::LoadedIndexDependentMemory; - descriptor.lowerability = simpleMixedLowerability( - "whole_module_pure_simt_requires_backend_check"); } else { return std::nullopt; } - descriptor.materializable = - compileOn91095 && - descriptor.lowerability.mixed == CandidateLoweringStatus::Native; + descriptor.materializable = compileOn91095 && descriptor.lowerability.mixed; return descriptor; } -static CandidateLoweringStatus -combineWholeKernelStatus(CandidateLoweringStatus lhs, - CandidateLoweringStatus rhs) { - auto rank = [](CandidateLoweringStatus status) { - switch (status) { - case CandidateLoweringStatus::Native: - return 0; - case CandidateLoweringStatus::BackendConditional: - return 1; - case CandidateLoweringStatus::AliasesMixed: - return 2; - case CandidateLoweringStatus::Unsupported: - return 3; - } - llvm_unreachable("unknown lowering status"); - }; - return rank(lhs) >= rank(rhs) ? lhs : rhs; -} - } // namespace llvm::StringRef mlir::ascend::stringifySimtAnchorKind(SimtAnchorKind kind) { @@ -671,21 +565,6 @@ llvm::StringRef mlir::ascend::stringifySimtAnchorKind(SimtAnchorKind kind) { llvm_unreachable("unknown SIMT anchor kind"); } -llvm::StringRef -mlir::ascend::stringifyCandidateLoweringStatus(CandidateLoweringStatus status) { - switch (status) { - case CandidateLoweringStatus::Unsupported: - return "unsupported"; - case CandidateLoweringStatus::Native: - return "native"; - case CandidateLoweringStatus::BackendConditional: - return "backend_conditional"; - case CandidateLoweringStatus::AliasesMixed: - return "aliases_mixed"; - } - llvm_unreachable("unknown candidate lowering status"); -} - llvm::SmallVector SimtAnchorPlan::materializableRoots() const { llvm::SmallVector result; result.reserve(anchors.size()); @@ -695,10 +574,72 @@ llvm::SmallVector SimtAnchorPlan::materializableRoots() const { return result; } -int64_t SimtAnchorPlan::materializableCount() const { - return llvm::count_if(anchors, [](const SimtAnchorDescriptor &anchor) { - return anchor.materializable; - }); +std::optional +mlir::ascend::mergeSimtStageAnchors(const SimtAnchorPlan &plan, + llvm::ArrayRef anchorIndices) { + llvm::SmallVector anchors; + llvm::DenseSet seen; + for (unsigned index : anchorIndices) { + if (index >= plan.anchors.size() || !seen.insert(index).second) + continue; + const SimtAnchorDescriptor &anchor = plan.anchors[index]; + if (!anchor.materializable || !anchor.operation) + return std::nullopt; + anchors.push_back(&anchor); + } + if (anchors.empty()) + return std::nullopt; + if (anchors.size() == 1) + return *anchors.front(); + + Block *block = nullptr; + llvm::DenseSet selected; + for (const SimtAnchorDescriptor *anchor : anchors) { + llvm::ArrayRef operations = anchor->scopeOperations; + if (operations.empty()) + operations = llvm::ArrayRef(anchor->operation); + for (Operation *operation : operations) { + if (!operation || !operation->getBlock() || + (block && block != operation->getBlock())) + return std::nullopt; + block = operation->getBlock(); + selected.insert(operation); + } + } + if (!block) + return std::nullopt; + + Operation *first = nullptr; + Operation *last = nullptr; + for (Operation &operation : *block) { + if (!selected.contains(&operation)) + continue; + if (!first) + first = &operation; + last = &operation; + } + if (!first || !last) + return std::nullopt; + + SimtAnchorDescriptor merged = *anchors.front(); + merged.scopeOperations.clear(); + bool inRange = false; + for (Operation &operation : *block) { + if (&operation == first) + inRange = true; + if (inRange) { + if (operation.hasTrait() || + operation.hasTrait() || + operation.getName().getStringRef() == "scope.scope" || + operation.getName().getStringRef() == "scope.return") + return std::nullopt; + merged.scopeOperations.push_back(&operation); + } + if (&operation == last) + break; + } + merged.scopeInsertionPoint = first; + return merged; } bool mlir::ascend::isLoadedIndexDependentMemoryOp(Operation *op) { @@ -709,13 +650,6 @@ bool mlir::ascend::isLoadedIndexDependentMemoryOp(Operation *op) { hasTensorPointerOperand(op) && pointerDependsOnLoadedIndex(op); } -std::optional -mlir::ascend::classifyMixedSimtAnchor(Operation *op) { - auto descriptor = analyzeAnchor(op, /*compileOn91095=*/true); - return descriptor ? std::optional(descriptor->kind) - : std::nullopt; -} - SimtAnchorPlan mlir::ascend::buildMixedSimtAnchorPlan(ModuleOp module, bool compileOn91095) { SimtAnchorPlan plan; @@ -733,9 +667,7 @@ SimtAnchorPlan mlir::ascend::buildMixedSimtAnchorPlan(ModuleOp module, op, descriptor->scopeInsertionPoint); if (descriptor->scopeOperations.empty()) { descriptor->materializable = false; - descriptor->lowerability.mixed = CandidateLoweringStatus::Unsupported; - descriptor->lowerability.mixedReasons.push_back( - "triangular_solve_has_no_materializable_scope_operations"); + descriptor->lowerability.mixed = false; } } else { descriptor->scopeOperations.push_back(op); @@ -747,43 +679,14 @@ SimtAnchorPlan mlir::ascend::buildMixedSimtAnchorPlan(ModuleOp module, return WalkResult::skip(); }); - bool anyMixedNative = false; + bool anyMixed = false; bool mixedBlocked = false; for (const SimtAnchorDescriptor &anchor : plan.anchors) { - plan.kernelLowerability.allSimd = combineWholeKernelStatus( - plan.kernelLowerability.allSimd, anchor.lowerability.allSimd); - plan.kernelLowerability.allSimtOnly = combineWholeKernelStatus( - plan.kernelLowerability.allSimtOnly, anchor.lowerability.allSimtOnly); - llvm::append_range(plan.kernelLowerability.allSimdReasons, - anchor.lowerability.allSimdReasons); - llvm::append_range(plan.kernelLowerability.allSimtOnlyReasons, - anchor.lowerability.allSimtOnlyReasons); - if (anchor.lowerability.mixed == CandidateLoweringStatus::Native) - anyMixedNative = true; - else if (anchor.lowerability.allSimd != CandidateLoweringStatus::Native) - mixedBlocked = true; - llvm::append_range(plan.kernelLowerability.mixedReasons, - anchor.lowerability.mixedReasons); - } - if (plan.anchors.empty()) { - plan.kernelLowerability.mixed = CandidateLoweringStatus::Unsupported; - plan.kernelLowerability.mixedReasons.push_back("no_recognized_simt_anchor"); - } else if (anyMixedNative && !mixedBlocked) { - plan.kernelLowerability.mixed = CandidateLoweringStatus::Native; - } else { - plan.kernelLowerability.mixed = - mixedBlocked ? CandidateLoweringStatus::Unsupported - : CandidateLoweringStatus::BackendConditional; + plan.kernelLowerability.allSimd &= anchor.lowerability.allSimd; + plan.kernelLowerability.allSimtOnly &= anchor.lowerability.allSimtOnly; + anyMixed |= anchor.lowerability.mixed; + mixedBlocked |= !anchor.lowerability.mixed && !anchor.lowerability.allSimd; } + plan.kernelLowerability.mixed = anyMixed && !mixedBlocked; return plan; } - -bool mlir::ascend::isMixedSimtAnchor(Operation *op, bool compileOn91095) { - auto descriptor = analyzeAnchor(op, compileOn91095); - return descriptor && descriptor->materializable; -} - -llvm::SmallVector -mlir::ascend::collectMixedSimtAnchors(ModuleOp module, bool compileOn91095) { - return buildMixedSimtAnchorPlan(module, compileOn91095).materializableRoots(); -} diff --git a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StagePartitioner.cpp b/third_party/ascend/costmodel/lib/AscendModel/Analysis/StagePartitioner.cpp similarity index 53% rename from third_party/ascend/costmodel/lib/AscendModel/RouteModel/StagePartitioner.cpp rename to third_party/ascend/costmodel/lib/AscendModel/Analysis/StagePartitioner.cpp index 6fd2b52d50..bc2b13fec7 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StagePartitioner.cpp +++ b/third_party/ascend/costmodel/lib/AscendModel/Analysis/StagePartitioner.cpp @@ -1,6 +1,6 @@ //===- StagePartitioner.cpp - Build semantic Phase/Stage IR -------------===// -#include "AscendModel/RouteModel/StagePartitioner.h" +#include "AscendModel/Analysis/StagePartitioner.h" #include "mlir/IR/BuiltinTypes.h" #include "llvm/ADT/DenseSet.h" @@ -21,41 +21,6 @@ using namespace mlir::ascend; namespace { -static double mapValue(const llvm::StringMap &map, - llvm::StringRef name) { - auto iterator = map.find(name); - return iterator == map.end() - ? 0.0 - : static_cast(std::max(0, iterator->second)); -} - -static StageWorkload buildWorkload(const SimtAnchorFeatureSummary &features) { - StageWorkload work; - const std::pair names[] = { - {"f32.add", "add"}, {"f32.sub", "sub"}, {"f32.mul", "mul"}, - {"f32.div", "div"}, {"f32.max", "max"}, {"f32.abs", "abs"}, - {"f32.exp", "exp"}, {"f32.log", "log"}, {"convert.cast", "cast"}, - {"f32.clamp", "clamp"}, - }; - for (const auto &[profileName, featureName] : names) { - const double elements = mapValue(features.opElements, featureName); - if (elements > 0.0) - work.operationElements[profileName] = elements; - } - work.scalarOperations = mapValue(features.weightedOps, "scalar"); - work.loadBytes = features.loadBytes; - work.storeBytes = features.storeBytes; - work.loadWarpInstructions = - static_cast(features.loadWarpInstructions); - work.storeWarpInstructions = - static_cast(features.storeWarpInstructions); - work.predicateElements = static_cast(std::max( - features.predicateLaneEvaluations, features.predicateElements)); - work.shuffleLaneSteps = static_cast(features.shuffleLaneSteps); - work.dotFlops = static_cast(features.dotFlops); - return work; -} - static void recomputeIssueElements(StageWorkload &work) { double elements = work.scalarOperations + work.predicateElements; for (const auto &entry : work.operationElements) @@ -348,98 +313,6 @@ static int64_t countAlgorithmLoops(const LogicalStage &stage) { return count; } -static double consume(double &remaining, double requested) { - const double value = - std::min(std::max(0.0, remaining), std::max(0.0, requested)); - remaining -= value; - return value; -} - -static StageWorkload consumeExact(StageWorkload &remaining, - const StageWorkload &requested) { - StageWorkload result; - result.scalarOperations = - consume(remaining.scalarOperations, requested.scalarOperations); - result.loadBytes = consume(remaining.loadBytes, requested.loadBytes); - result.storeBytes = consume(remaining.storeBytes, requested.storeBytes); - result.loadWarpInstructions = - consume(remaining.loadWarpInstructions, requested.loadWarpInstructions); - result.storeWarpInstructions = - consume(remaining.storeWarpInstructions, requested.storeWarpInstructions); - result.predicateElements = - consume(remaining.predicateElements, requested.predicateElements); - result.shuffleLaneSteps = - consume(remaining.shuffleLaneSteps, requested.shuffleLaneSteps); - result.dotFlops = consume(remaining.dotFlops, requested.dotFlops); - result.estimatedSpillTransactions = - consume(remaining.estimatedSpillTransactions, - requested.estimatedSpillTransactions); - for (const auto &[name, elements] : requested.operationElements) { - double &available = remaining.operationElements[name]; - const double owned = consume(available, elements); - if (owned > 0.0) - result.operationElements[name] = owned; - } - recomputeIssueElements(result); - recomputeIssueElements(remaining); - return result; -} - -static StageWorkload takeScalarAndPredicate(StageWorkload &remaining) { - StageWorkload result; - result.scalarOperations = std::exchange(remaining.scalarOperations, 0.0); - result.predicateElements = std::exchange(remaining.predicateElements, 0.0); - recomputeIssueElements(result); - recomputeIssueElements(remaining); - return result; -} - -static StageWorkload takeLoads(StageWorkload &remaining) { - StageWorkload result; - result.loadBytes = std::exchange(remaining.loadBytes, 0.0); - result.loadWarpInstructions = - std::exchange(remaining.loadWarpInstructions, 0.0); - recomputeIssueElements(result); - recomputeIssueElements(remaining); - return result; -} - -static StageWorkload takeStores(StageWorkload &remaining) { - StageWorkload result; - result.storeBytes = std::exchange(remaining.storeBytes, 0.0); - result.storeWarpInstructions = - std::exchange(remaining.storeWarpInstructions, 0.0); - recomputeIssueElements(result); - recomputeIssueElements(remaining); - return result; -} - -static StageWorkload takeDot(StageWorkload &remaining) { - StageWorkload result; - result.dotFlops = std::exchange(remaining.dotFlops, 0.0); - recomputeIssueElements(result); - recomputeIssueElements(remaining); - return result; -} - -static void moveOperation(StageWorkload &from, StageWorkload &to, - llvm::StringRef name) { - auto iterator = from.operationElements.find(name); - if (iterator == from.operationElements.end()) - return; - to.operationElements[name] += iterator->second; - from.operationElements.erase(iterator); -} - -static StageWorkload takeAllOperations(StageWorkload &remaining) { - StageWorkload result; - result.operationElements = std::move(remaining.operationElements); - remaining.operationElements.clear(); - recomputeIssueElements(result); - recomputeIssueElements(remaining); - return result; -} - static void mergeWorkload(StageWorkload &into, StageWorkload from) { into.scalarOperations += from.scalarOperations; into.loadBytes += from.loadBytes; @@ -473,13 +346,11 @@ static void makePerIteration(LogicalStage &stage) { recomputeIssueElements(work); } -static LogicalStage makeStage(llvm::StringRef id, llvm::StringRef description, - StageCostModelKind kind, +static LogicalStage makeStage(llvm::StringRef id, StageCostModelKind kind, StageScheduleKind schedule, int64_t iterations, StageWorkload workload) { LogicalStage stage; stage.id = id.str(); - stage.description = description.str(); stage.costModelKind = kind; stage.scheduleKind = schedule; stage.iterationCount = std::max(1, iterations); @@ -505,254 +376,148 @@ static LogicalStage asLocalSIMT(LogicalStage stage) { } static void addPhase(StagePartition &partition, llvm::StringRef id, - llvm::StringRef description, LogicalStage stage) { + LogicalStage stage) { LogicalPhase phase; phase.id = id.str(); - phase.description = description.str(); phase.stages.push_back(std::move(stage)); partition.phases.push_back(std::move(phase)); } static bool operationTreeContainsName(Operation *root, llvm::StringRef name); -static bool hasWork(const StageWorkload &work) { - if (work.scalarOperations > 0.0 || work.loadBytes > 0.0 || - work.storeBytes > 0.0 || work.predicateElements > 0.0 || - work.shuffleLaneSteps > 0.0 || work.dotFlops > 0.0 || - work.estimatedSpillTransactions > 0.0) - return true; - return llvm::any_of(work.operationElements, - [](const auto &entry) { return entry.second > 0.0; }); +static bool hasPhase(const PhaseBoundaryPlan *plan, llvm::StringRef id) { + return plan && llvm::is_contained(plan->rootPhaseIds, id); } static void prependAutoBlockifyStages(StagePartition &partition, - StageWorkload &remaining, const SimdSimtFeatureSummary &features, - bool operationGraphHasAutoBlockify) { - if (!features.autoBlockifyV1Applied && !operationGraphHasAutoBlockify) + const PhaseBoundaryPlan *plan) { + if (!features.autoBlockifyV1Applied && + !hasPhase(plan, "auto_blockify_dispatch")) return; - StageWorkload dispatch; - dispatch.scalarOperations = - consume(remaining.scalarOperations, - static_cast(features.autoBlockifyV1ScheduleOpCount)); - recomputeIssueElements(dispatch); LogicalPhase phase; phase.id = "auto_blockify_dispatch"; - phase.description = "AutoBlockify V1 physical/logical program dispatch"; - phase.stages.push_back(makeStage( - "physical_program_dispatch", "Map physical PID to logical block range", - StageCostModelKind::AutoBlockifyDispatch, StageScheduleKind::StraightLine, - 1, std::move(dispatch))); - if (features.autoBlockifyV1LoopCount > 0 || operationGraphHasAutoBlockify) { - StageWorkload loop; - loop.scalarOperations = - consume(remaining.scalarOperations, - static_cast(features.autoBlockifyV1LoopCount)); - recomputeIssueElements(loop); - phase.stages.push_back(makeStage( - "logical_program_loop", "Iterate AutoBlockify V1 logical programs", - StageCostModelKind::AutoBlockifyLoop, - StageScheduleKind::IndependentPipelined, - std::max(1, features.autoBlockifyV1LoopCount), - std::move(loop))); - } + phase.stages.push_back(makeStage("physical_program_dispatch", + StageCostModelKind::AutoBlockifyDispatch, + StageScheduleKind::StraightLine, 1, {})); + if (features.autoBlockifyV1LoopCount > 0 || + hasPhase(plan, "auto_blockify_dispatch")) + phase.stages.push_back( + makeStage("logical_program_loop", StageCostModelKind::AutoBlockifyLoop, + StageScheduleKind::IndependentPipelined, + std::max(1, features.autoBlockifyV1LoopCount), {})); partition.phases.push_back(std::move(phase)); } static StagePartition partitionTriangular(const SimdSimtFeatureSummary &features, const TriangularSolveFacts &facts, - const PhaseBoundaryPlan *operationGraphPlan) { + const PhaseBoundaryPlan *plan) { StagePartition partition; partition.domain = "triangular_recurrence"; - StageWorkload remaining = buildKernelStageWorkload(features); - StageWorkload anchor = buildWorkload(features.simtAnchors); - anchor = consumeExact(remaining, anchor); - - const bool graphHasAutoBlockify = - operationGraphPlan && llvm::is_contained(operationGraphPlan->rootPhaseIds, - "auto_blockify_dispatch"); - prependAutoBlockifyStages(partition, remaining, features, - graphHasAutoBlockify); - - StageWorkload head = takeScalarAndPredicate(remaining); - mergeWorkload(head, takeAllOperations(remaining)); - head.paysKernelSetup = true; - addPhase(partition, "head", "Tile offset and triangular mask setup", - withControl(makeStage("head_index_mask", - "Scalar indices and triangular masks", - StageCostModelKind::PredicateMask, - StageScheduleKind::StraightLine, 1, - std::move(head)), - features.conditionalBranchCount - - features.simtAnchors.conditionalBranchCount, - features.divergentBranchCount - - features.simtAnchors.divergentBranchCount, - features.activeLaneRatio)); - - StageWorkload loads = takeLoads(remaining); - const bool graphHasDiagonalLoad = - operationGraphPlan && - llvm::is_contained(operationGraphPlan->rootPhaseIds, "diagonal_load"); - if (hasWork(loads) || graphHasDiagonalLoad) - addPhase(partition, "diagonal_load", "Load diagonal tile data", - makeStage("load_diagonal_tiles", "Continuous diagonal tile loads", + prependAutoBlockifyStages(partition, features, plan); + + LogicalStage head = withControl( + makeStage("head_index_mask", StageCostModelKind::PredicateMask, + StageScheduleKind::StraightLine, 1, {}), + features.conditionalBranchCount - + features.simtAnchors.conditionalBranchCount, + features.divergentBranchCount - features.simtAnchors.divergentBranchCount, + features.activeLaneRatio); + head.workload.paysKernelSetup = true; + addPhase(partition, "head", std::move(head)); + + if (hasPhase(plan, "diagonal_load")) + addPhase(partition, "diagonal_load", + makeStage("load_diagonal_tiles", StageCostModelKind::ContinuousTileMemory, - StageScheduleKind::IndependentPipelined, 1, - std::move(loads))); - - const int64_t recurrenceIterations = - std::max(1, facts.recurrenceLoopCount); - // The feature-summary fallback has no owned operation graph from which to - // count sibling loops. Derive the number of independent 16x16 recurrence - // groups from the structural triangular facts. Production operation-graph - // analysis below recomputes the same field directly from owned loops. + StageScheduleKind::IndependentPipelined, 1, {})); + + const int64_t iterations = std::max(1, facts.recurrenceLoopCount); LogicalStage recurrence = asLocalSIMT(withControl( makeStage("diagonal_inverse_recurrence", - "Predicate, short reduction and recurrent state update", StageCostModelKind::LoopCarriedRecurrence, - StageScheduleKind::LoopCarriedSerial, recurrenceIterations, - std::move(anchor)), + StageScheduleKind::LoopCarriedSerial, iterations, {}), features.simtAnchors.conditionalBranchCount, features.simtAnchors.divergentBranchCount, features.simtAnchors.activeLaneRatio)); - const int64_t iterationsPerGroup = + const int64_t rows = std::max(1, facts.blockRows - facts.recurrenceStartRow); - recurrence.features.parallelRecurrenceGroupCount = std::max( - 1, (recurrenceIterations + iterationsPerGroup - 1) / iterationsPerGroup); - addPhase(partition, "diagonal_inverse", - "Loop-carried blockwise triangular recurrence", - std::move(recurrence)); - - StageWorkload dot = takeDot(remaining); - StageWorkload stores = takeStores(remaining); - mergeWorkload(stores, std::move(remaining)); - bool graphHasDot = false; - bool graphHasStore = false; - if (operationGraphPlan) - for (auto indexedRoot : - llvm::enumerate(operationGraphPlan->rootOperations)) { - if (operationGraphPlan->rootPhaseIds[indexedRoot.index()] != - "merge_store") + recurrence.features.parallelRecurrenceGroupCount = + std::max(1, (iterations + rows - 1) / rows); + addPhase(partition, "diagonal_inverse", std::move(recurrence)); + + if (hasPhase(plan, "merge_store")) { + bool hasDot = false; + bool hasStore = false; + for (auto root : llvm::enumerate(plan->rootOperations)) { + if (plan->rootPhaseIds[root.index()] != "merge_store") continue; - graphHasDot |= operationTreeContainsName(indexedRoot.value(), "tt.dot"); - graphHasStore |= - operationTreeContainsName(indexedRoot.value(), "tt.store"); + hasDot |= operationTreeContainsName(root.value(), "tt.dot"); + hasStore |= operationTreeContainsName(root.value(), "tt.store"); } - if (hasWork(dot) || hasWork(stores) || graphHasDot || graphHasStore) { - LogicalPhase mergeStore; - mergeStore.id = "merge_store"; - mergeStore.description = "Dense off-diagonal merge and result store"; - if (hasWork(dot) || graphHasDot) - mergeStore.stages.push_back(makeStage( - "dense_dot_tail", "Dense dot tail preserved for Cube", - StageCostModelKind::CubeRoofline, - StageScheduleKind::IndependentPipelined, - std::max(1, facts.denseDotTailOps), std::move(dot))); - if (hasWork(stores) || graphHasStore) - mergeStore.stages.push_back( - makeStage("store_inverse_tile", "Continuous result tile store", - StageCostModelKind::ContinuousTileStore, - StageScheduleKind::StraightLine, 1, std::move(stores))); - partition.phases.push_back(std::move(mergeStore)); + LogicalPhase phase; + phase.id = "merge_store"; + if (hasDot) + phase.stages.push_back( + makeStage("dense_dot_tail", StageCostModelKind::CubeRoofline, + StageScheduleKind::IndependentPipelined, + std::max(1, facts.denseDotTailOps), {})); + if (hasStore) + phase.stages.push_back(makeStage("store_inverse_tile", + StageCostModelKind::ContinuousTileStore, + StageScheduleKind::StraightLine, 1, {})); + if (!phase.stages.empty()) + partition.phases.push_back(std::move(phase)); } return partition; } -static StagePartition -partitionRowwise(const SimdSimtFeatureSummary &features, - const PhaseBoundaryPlan *operationGraphPlan) { +static StagePartition partitionRowwise(const SimdSimtFeatureSummary &features, + const PhaseBoundaryPlan *plan) { StagePartition partition; partition.domain = "loaded_index_rowwise_reduction"; - StageWorkload remaining = buildKernelStageWorkload(features); - const bool graphHasAutoBlockify = - operationGraphPlan && llvm::is_contained(operationGraphPlan->rootPhaseIds, - "auto_blockify_dispatch"); - prependAutoBlockifyStages(partition, remaining, features, - graphHasAutoBlockify); - - StageWorkload index = takeScalarAndPredicate(remaining); - index.paysKernelSetup = true; - addPhase(partition, "row_dispatch", "Resolve token and row indices", - withControl( - makeStage("row_index_generation", "Scalar row index and masks", - StageCostModelKind::IndexGeneration, - StageScheduleKind::StraightLine, 1, std::move(index)), - features.conditionalBranchCount, features.divergentBranchCount, - features.activeLaneRatio)); - - addPhase(partition, "row_load", "Loaded-index-dependent row access", - asLocalSIMT(makeStage("indirect_row_gather", - "Gather the selected row tile", - StageCostModelKind::IndirectGatherMemory, - StageScheduleKind::PartiallyDependent, 1, - takeLoads(remaining)))); - - StageWorkload reduction; - reduction.shuffleLaneSteps = std::exchange(remaining.shuffleLaneSteps, 0.0); - moveOperation(remaining, reduction, "f32.max"); - recomputeIssueElements(reduction); - const int64_t reductionIterations = + prependAutoBlockifyStages(partition, features, plan); + addPhase(partition, "row_dispatch", + makeStage("row_index_generation", + StageCostModelKind::IndexGeneration, + StageScheduleKind::StraightLine, 1, {})); + addPhase(partition, "row_load", + asLocalSIMT(makeStage( + "indirect_row_gather", StageCostModelKind::IndirectGatherMemory, + StageScheduleKind::PartiallyDependent, 1, {}))); + const int64_t iterations = std::max(1, features.staticLoopTripCountMax); - addPhase(partition, "row_reduction", "Reduce each selected row", - makeStage("rowwise_reduction", "Row-local reduction tree", - StageCostModelKind::RowwiseReduction, - StageScheduleKind::PartiallyDependent, reductionIterations, - std::move(reduction))); - - StageWorkload convert = takeStores(remaining); - mergeWorkload(convert, takeAllOperations(remaining)); - mergeWorkload(convert, std::move(remaining)); - addPhase(partition, "convert_store", "Scale, convert, pack and store", - makeStage("conversion_pack_store", "Conversion and packed output", + addPhase(partition, "row_reduction", + makeStage("rowwise_reduction", StageCostModelKind::RowwiseReduction, + StageScheduleKind::PartiallyDependent, iterations, {})); + addPhase(partition, "convert_store", + makeStage("conversion_pack_store", StageCostModelKind::ConversionPack, - StageScheduleKind::IndependentPipelined, - reductionIterations, std::move(convert))); + StageScheduleKind::IndependentPipelined, iterations, {})); return partition; } static StagePartition partitionIndirectDot(const SimdSimtFeatureSummary &features, - const PhaseBoundaryPlan *operationGraphPlan) { + const PhaseBoundaryPlan *plan) { StagePartition partition; partition.domain = "indirect_underfilled_dot"; - StageWorkload remaining = buildKernelStageWorkload(features); - const bool graphHasAutoBlockify = - operationGraphPlan && llvm::is_contained(operationGraphPlan->rootPhaseIds, - "auto_blockify_dispatch"); - prependAutoBlockifyStages(partition, remaining, features, - graphHasAutoBlockify); - - StageWorkload index = takeScalarAndPredicate(remaining); - mergeWorkload(index, takeAllOperations(remaining)); - index.paysKernelSetup = true; - addPhase(partition, "index_setup", "Generate gather indices and masks", - withControl( - makeStage("index_generation", "Index and predicate generation", - StageCostModelKind::IndexGeneration, - StageScheduleKind::StraightLine, 1, std::move(index)), - features.conditionalBranchCount, features.divergentBranchCount, - features.activeLaneRatio)); - - addPhase(partition, "gather_tiles", "Gather dot input tiles", - asLocalSIMT(makeStage("indirect_tile_gather", - "Loaded-index-dependent operand gathers", - StageCostModelKind::IndirectGatherMemory, - StageScheduleKind::PartiallyDependent, 1, - takeLoads(remaining)))); - - addPhase(partition, "dot", "Under-filled matrix product", - makeStage("tiny_cube_dot", "Small dot with Cube underfill", - StageCostModelKind::TinyCubeRoofline, - StageScheduleKind::IndependentPipelined, 1, - takeDot(remaining))); - - StageWorkload store = takeStores(remaining); - mergeWorkload(store, std::move(remaining)); - addPhase(partition, "output_store", "Write dot result", - makeStage("store_dot_result", "Continuous result store", + prependAutoBlockifyStages(partition, features, plan); + addPhase(partition, "index_setup", + makeStage("index_generation", StageCostModelKind::IndexGeneration, + StageScheduleKind::StraightLine, 1, {})); + addPhase(partition, "gather_tiles", + asLocalSIMT(makeStage( + "indirect_tile_gather", StageCostModelKind::IndirectGatherMemory, + StageScheduleKind::PartiallyDependent, 1, {}))); + addPhase(partition, "dot", + makeStage("tiny_cube_dot", StageCostModelKind::TinyCubeRoofline, + StageScheduleKind::IndependentPipelined, 1, {})); + addPhase(partition, "output_store", + makeStage("store_dot_result", StageCostModelKind::ContinuousTileStore, - StageScheduleKind::StraightLine, 1, std::move(store))); + StageScheduleKind::StraightLine, 1, {})); return partition; } @@ -771,38 +536,6 @@ static bool anchorMatchesStage(const SimtAnchorDescriptor &anchor, static Operation *getTopLevelSemanticRoot(Operation *operation); -static void attachAnchorOperationOwnership(StagePartition &partition, - const SimtAnchorPlan &anchorPlan) { - for (LogicalPhase &phase : partition.phases) { - for (LogicalStage &stage : phase.stages) { - if (!stage.localSimtMaterializable) - continue; - for (auto indexedAnchor : llvm::enumerate(anchorPlan.anchors)) { - const SimtAnchorDescriptor &anchor = indexedAnchor.value(); - if (!anchorMatchesStage(anchor, stage)) - continue; - stage.simtAnchorIndices.push_back( - static_cast(indexedAnchor.index())); - if (anchor.scopeOperations.empty()) { - if (anchor.operation && - !llvm::is_contained(stage.operations, anchor.operation)) - stage.operations.push_back(anchor.operation); - continue; - } - for (Operation *operation : anchor.scopeOperations) - if (operation && !llvm::is_contained(stage.operations, operation)) - stage.operations.push_back(operation); - } - // A production mixed Stage must be backed by the exact operations that - // the materializer will consume. Synthetic feature-only tests use the - // overload without an anchor plan and retain fallback behavior. - stage.localSimtMaterializable = !stage.operations.empty(); - if (!stage.localSimtMaterializable) - stage.localSimtFactors.clear(); - } - } -} - static bool stageOwnsAnchor(const LogicalStage &stage, const SimtAnchorDescriptor &anchor) { if (!anchorMatchesStage(anchor, stage)) @@ -915,12 +648,6 @@ static bool operationTreeContainsLoadedIndexMemory(Operation *root) { /// monotone: after a boundary is crossed, a later root cannot move back to an /// earlier Phase. Cost and candidate mode are intentionally absent here. static llvm::Error assignRootPhaseIds(PhaseBoundaryPlan &plan) { - enum class TriangularPhase { Head, Load, Recurrence, MergeStore }; - enum class RowwisePhase { Index, Gather, Reduction, ConvertStore }; - enum class IndirectDotPhase { Index, Gather, Dot, OutputStore }; - TriangularPhase triangular = TriangularPhase::Head; - RowwisePhase rowwise = RowwisePhase::Index; - IndirectDotPhase indirectDot = IndirectDotPhase::Index; llvm::DenseSet anchorRoots(plan.localSimtAnchorRoots.begin(), plan.localSimtAnchorRoots.end()); std::optional firstAnchorIndex; @@ -949,6 +676,7 @@ static llvm::Error assignRootPhaseIds(PhaseBoundaryPlan &plan) { plan.rootPhaseIds.clear(); plan.rootPhaseIds.reserve(plan.rootOperations.size()); + llvm::StringRef current; for (auto indexedRoot : llvm::enumerate(plan.rootOperations)) { Operation *root = indexedRoot.value(); if (!root) @@ -965,79 +693,43 @@ static llvm::Error assignRootPhaseIds(PhaseBoundaryPlan &plan) { case PhaseBoundaryDomain::TriangularRecurrence: if (firstAnchorIndex && indexedRoot.index() >= *firstAnchorIndex && indexedRoot.index() <= *lastAnchorIndex) - triangular = TriangularPhase::Recurrence; + current = "diagonal_inverse"; else if (lastAnchorIndex && indexedRoot.index() > *lastAnchorIndex) - triangular = TriangularPhase::MergeStore; + current = "merge_store"; else if (operationTreeContainsName(root, "tt.dot") || operationTreeContainsName(root, "tt.store")) - triangular = TriangularPhase::MergeStore; - else if (triangular == TriangularPhase::Head && + current = "merge_store"; + else if ((current.empty() || current == "head") && operationTreeContainsName(root, "tt.load")) - triangular = TriangularPhase::Load; - switch (triangular) { - case TriangularPhase::Head: - plan.rootPhaseIds.push_back("head"); - break; - case TriangularPhase::Load: - plan.rootPhaseIds.push_back("diagonal_load"); - break; - case TriangularPhase::Recurrence: - plan.rootPhaseIds.push_back("diagonal_inverse"); - break; - case TriangularPhase::MergeStore: - plan.rootPhaseIds.push_back("merge_store"); - break; - } + current = "diagonal_load"; + else if (current.empty()) + current = "head"; break; case PhaseBoundaryDomain::LoadedIndexRowwiseReduction: if (operationTreeContainsName(root, "tt.reduce")) - rowwise = RowwisePhase::Reduction; - else if (rowwise == RowwisePhase::Reduction || + current = "row_reduction"; + else if (current == "row_reduction" || operationTreeContainsName(root, "tt.store")) - rowwise = RowwisePhase::ConvertStore; - else if (rowwise == RowwisePhase::Index && + current = "convert_store"; + else if ((current.empty() || current == "row_dispatch") && operationTreeContainsLoadedIndexMemory(root)) - rowwise = RowwisePhase::Gather; - switch (rowwise) { - case RowwisePhase::Index: - plan.rootPhaseIds.push_back("row_dispatch"); - break; - case RowwisePhase::Gather: - plan.rootPhaseIds.push_back("row_load"); - break; - case RowwisePhase::Reduction: - plan.rootPhaseIds.push_back("row_reduction"); - break; - case RowwisePhase::ConvertStore: - plan.rootPhaseIds.push_back("convert_store"); - break; - } + current = "row_load"; + else if (current.empty()) + current = "row_dispatch"; break; case PhaseBoundaryDomain::IndirectUnderfilledDot: if (operationTreeContainsName(root, "tt.dot")) - indirectDot = IndirectDotPhase::Dot; - else if (indirectDot == IndirectDotPhase::Dot || - operationTreeContainsName(root, "tt.store")) - indirectDot = IndirectDotPhase::OutputStore; - else if (indirectDot == IndirectDotPhase::Index && + current = "dot"; + else if (current == "dot" || operationTreeContainsName(root, "tt.store")) + current = "output_store"; + else if ((current.empty() || current == "index_setup") && operationTreeContainsLoadedIndexMemory(root)) - indirectDot = IndirectDotPhase::Gather; - switch (indirectDot) { - case IndirectDotPhase::Index: - plan.rootPhaseIds.push_back("index_setup"); - break; - case IndirectDotPhase::Gather: - plan.rootPhaseIds.push_back("gather_tiles"); - break; - case IndirectDotPhase::Dot: - plan.rootPhaseIds.push_back("dot"); - break; - case IndirectDotPhase::OutputStore: - plan.rootPhaseIds.push_back("output_store"); - break; - } + current = "gather_tiles"; + else if (current.empty()) + current = "index_setup"; break; } + plan.rootPhaseIds.push_back(current.str()); } if (plan.rootPhaseIds.size() != plan.rootOperations.size()) return llvm::createStringError( @@ -1068,6 +760,29 @@ static LogicalStage *findStage(StagePartition &partition, llvm::StringRef id) { return nullptr; } +static llvm::StringRef stageIdForPhase(PhaseBoundaryDomain domain, + llvm::StringRef phaseId) { + if (domain == PhaseBoundaryDomain::TriangularRecurrence) + return llvm::StringSwitch(phaseId) + .Case("head", "head_index_mask") + .Case("diagonal_load", "load_diagonal_tiles") + .Case("diagonal_inverse", "diagonal_inverse_recurrence") + .Default({}); + if (domain == PhaseBoundaryDomain::LoadedIndexRowwiseReduction) + return llvm::StringSwitch(phaseId) + .Case("row_dispatch", "row_index_generation") + .Case("row_load", "indirect_row_gather") + .Case("row_reduction", "rowwise_reduction") + .Case("convert_store", "conversion_pack_store") + .Default({}); + return llvm::StringSwitch(phaseId) + .Case("index_setup", "index_generation") + .Case("gather_tiles", "indirect_tile_gather") + .Case("dot", "tiny_cube_dot") + .Case("output_store", "store_dot_result") + .Default({}); +} + static int64_t getStageOrdinal(const StagePartition &partition, const LogicalStage *target) { int64_t ordinal = 0; @@ -1125,44 +840,14 @@ attachCompleteOperationOwnership(StagePartition &partition, : "physical_program_dispatch"); } - if (!target) { - switch (plan.domain) { - case PhaseBoundaryDomain::TriangularRecurrence: - if (phaseId == "head") - target = findStage(partition, "head_index_mask"); - else if (phaseId == "diagonal_load") - target = findStage(partition, "load_diagonal_tiles"); - else if (phaseId == "diagonal_inverse") - target = findStage(partition, "diagonal_inverse_recurrence"); - else if (phaseId == "merge_store") { - if (operationTreeContainsName(root, "tt.store")) - mergeStoreReached = true; - target = findStage(partition, mergeStoreReached ? "store_inverse_tile" - : "dense_dot_tail"); - } - break; - case PhaseBoundaryDomain::LoadedIndexRowwiseReduction: - if (phaseId == "row_dispatch") - target = findStage(partition, "row_index_generation"); - else if (phaseId == "row_load") - target = findStage(partition, "indirect_row_gather"); - else if (phaseId == "row_reduction") - target = findStage(partition, "rowwise_reduction"); - else if (phaseId == "convert_store") - target = findStage(partition, "conversion_pack_store"); - break; - case PhaseBoundaryDomain::IndirectUnderfilledDot: - if (phaseId == "index_setup") - target = findStage(partition, "index_generation"); - else if (phaseId == "gather_tiles") - target = findStage(partition, "indirect_tile_gather"); - else if (phaseId == "dot") - target = findStage(partition, "tiny_cube_dot"); - else if (phaseId == "output_store") - target = findStage(partition, "store_dot_result"); - break; - } + if (!target && plan.domain == PhaseBoundaryDomain::TriangularRecurrence && + phaseId == "merge_store") { + mergeStoreReached |= operationTreeContainsName(root, "tt.store"); + target = findStage(partition, mergeStoreReached ? "store_inverse_tile" + : "dense_dot_tail"); } + if (!target) + target = findStage(partition, stageIdForPhase(plan.domain, phaseId)); const int64_t ordinal = getStageOrdinal(partition, target); if (ordinal < 0) return llvm::createStringError( @@ -1183,7 +868,6 @@ attachCompleteOperationOwnership(StagePartition &partition, return llvm::createStringError( std::errc::invalid_argument, "StageBoundaryAnalysis did not conserve TTIR operation ownership"); - partition.boundarySource = "operation_graph"; partition.operationOwnershipComplete = true; partition.modeledOperationCount = static_cast(plan.rootOperations.size()); @@ -1284,25 +968,17 @@ static void deriveLocalSimtScopeTraffic(StagePartition &partition, stage.localSimtScopeCount = 0; stage.scopeInputTensorBytes = 0; stage.scopeOutputTensorBytes = 0; - llvm::DenseSet coveredByRange; - - for (unsigned anchorIndex : stage.simtAnchorIndices) { - if (anchorIndex >= anchorPlan.anchors.size()) - continue; - const SimtAnchorDescriptor &anchor = anchorPlan.anchors[anchorIndex]; - if (!anchor.materializable || !anchor.operation || - coveredByRange.contains(anchor.operation)) - continue; - + auto merged = mergeSimtStageAnchors(anchorPlan, stage.simtAnchorIndices); + if (!merged) + continue; + { + const SimtAnchorDescriptor &anchor = *merged; llvm::SmallVector roots; const bool isRange = anchor.scopeOperations.size() > 1; - if (isRange) { + if (isRange) llvm::append_range(roots, anchor.scopeOperations); - for (Operation *operation : roots) - coveredByRange.insert(operation); - } else { + else roots.push_back(anchor.operation); - } llvm::DenseSet inside; for (Operation *root : roots) { @@ -1333,6 +1009,21 @@ static void deriveLocalSimtScopeTraffic(StagePartition &partition, } } + // TritonToUnstructure cannot reconstruct offset information for a + // tensor-of-pointer returned by scope.scope. Capturing pointers is + // legal (the scope is not isolated from above), but returning pointer + // state would make this local Mixed implementation fail after route + // selection. Reject that implementation before it is scored; the + // same Stage remains legal in a whole-kernel pure-SIMT route. + if (llvm::any_of(returned, [](Value value) { + return isPointerLikeType(value.getType()); + })) { + stage.localSimtMaterializable = false; + stage.localSimtFactors.clear(); + stage.simtAnchorIndices.clear(); + continue; + } + ++stage.localSimtScopeCount; stage.scopeInputTensorBytes += staticTensorBytes(captured.getArrayRef()); @@ -1343,62 +1034,8 @@ static void deriveLocalSimtScopeTraffic(StagePartition &partition, } } -static double totalOperationElements(const StageWorkload &work) { - double result = 0.0; - for (const auto &entry : work.operationElements) - result += entry.second; - return result; -} - -static StageWorkload accumulatePartition(const StagePartition &partition) { - StageWorkload total; - for (const LogicalPhase &phase : partition.phases) { - for (const LogicalStage &stage : phase.stages) { - StageWorkload work = stage.workload; - const double count = - static_cast(std::max(1, stage.iterationCount)); - work.scalarOperations *= count; - work.loadBytes *= count; - work.storeBytes *= count; - work.loadWarpInstructions *= count; - work.storeWarpInstructions *= count; - work.predicateElements *= count; - work.shuffleLaneSteps *= count; - work.dotFlops *= count; - work.estimatedSpillTransactions *= count; - for (auto &entry : work.operationElements) - entry.second *= count; - mergeWorkload(total, std::move(work)); - } - } - return total; -} - -static bool near(double lhs, double rhs) { - return std::abs(lhs - rhs) <= 1.0e-6 * std::max({1.0, lhs, rhs}); -} - } // namespace -StageWorkload -mlir::ascend::buildKernelStageWorkload(const SimdSimtFeatureSummary &features) { - SimtAnchorFeatureSummary kernel; - kernel.opElements = features.opElements; - kernel.weightedOps = features.weightedOps; - kernel.loadBytes = features.loadBytes; - kernel.storeBytes = features.storeBytes; - kernel.loadWarpInstructions = features.loadWarpInstructions; - kernel.storeWarpInstructions = features.storeWarpInstructions; - kernel.predicateElements = features.predicateElements; - kernel.predicateLaneEvaluations = features.predicateLaneEvaluations; - kernel.shuffleLaneSteps = features.shuffleLaneSteps; - kernel.dotFlops = features.dotFlops; - StageWorkload work = buildWorkload(kernel); - work.scalarOperations = static_cast(features.scalarOps); - recomputeIssueElements(work); - return work; -} - llvm::Expected ProgramStructureAnalysis::analyze(ModuleOp module, const SimtAnchorPlan &anchorPlan) const { @@ -1474,22 +1111,22 @@ ProgramStructureAnalysis::analyze(ModuleOp module, return structure; } -llvm::Expected> -PhaseBoundaryAnalysis::analyze(const SimdSimtFeatureSummary &features, - const StagePartitionerOptions &options) const { +static std::optional +identifyPhaseBoundary(const SimdSimtFeatureSummary &features, + const StagePartitionerOptions &options) { if (features.simtAnchors.triangularSolves.size() == 1 && features.simtAnchors.count > 0) { PhaseBoundaryPlan plan{PhaseBoundaryDomain::TriangularRecurrence, "triangular_recurrence", features.simtAnchors.triangularSolves.front()}; - return std::optional{std::move(plan)}; + return plan; } if (features.dotOps == 0 && features.reduceOps > 0 && features.loadedIndexDependentMemoryOps > 0 && features.loadOps > 0 && features.storeOps > 0) { PhaseBoundaryPlan plan{PhaseBoundaryDomain::LoadedIndexRowwiseReduction, "loaded_index_rowwise_reduction", std::nullopt}; - return std::optional{std::move(plan)}; + return plan; } if (features.dotOps > 0 && features.reduceOps == 0 && features.loadedIndexDependentMemoryOps > 0 && @@ -1497,7 +1134,7 @@ PhaseBoundaryAnalysis::analyze(const SimdSimtFeatureSummary &features, features.storeOps > 0) { PhaseBoundaryPlan plan{PhaseBoundaryDomain::IndirectUnderfilledDot, "indirect_underfilled_dot", std::nullopt}; - return std::optional{std::move(plan)}; + return plan; } return std::optional{}; } @@ -1507,25 +1144,27 @@ PhaseBoundaryAnalysis::analyze(ModuleOp module, const SimtAnchorPlan &anchorPlan, const SimdSimtFeatureSummary &features, const StagePartitionerOptions &options) const { - auto plan = analyze(features, options); + auto plan = identifyPhaseBoundary(features, options); if (!plan) - return plan.takeError(); - if (!*plan) - return plan; + return std::optional{}; auto structure = ProgramStructureAnalysis().analyze(module, anchorPlan); if (!structure) return structure.takeError(); - (*plan)->rootOperations = std::move(structure->rootOperations); - (*plan)->localSimtAnchorRoots = std::move(structure->localSimtAnchorRoots); - if (llvm::Error error = assignRootPhaseIds(**plan)) + plan->rootOperations = std::move(structure->rootOperations); + plan->localSimtAnchorRoots = std::move(structure->localSimtAnchorRoots); + if (llvm::Error error = assignRootPhaseIds(*plan)) return std::move(error); - return plan; + return std::optional{std::move(*plan)}; } llvm::Expected StageBoundaryAnalysis::analyze(const PhaseBoundaryPlan &phasePlan, const SimdSimtFeatureSummary &features, const SimtAnchorPlan *anchorPlan) const { + if (!phasePlan.hasOperationGraph() || !anchorPlan) + return llvm::createStringError( + std::errc::invalid_argument, + "StageBoundaryAnalysis requires PreparedTTIR ownership"); StagePartition partition; switch (phasePlan.domain) { case PhaseBoundaryDomain::TriangularRecurrence: @@ -1533,32 +1172,23 @@ StageBoundaryAnalysis::analyze(const PhaseBoundaryPlan &phasePlan, return llvm::createStringError( std::errc::invalid_argument, "triangular PhaseBoundaryPlan has no recurrence facts"); - partition = partitionTriangular(features, *phasePlan.triangularSolve, - phasePlan.hasOperationGraph() ? &phasePlan - : nullptr); + partition = + partitionTriangular(features, *phasePlan.triangularSolve, &phasePlan); break; case PhaseBoundaryDomain::LoadedIndexRowwiseReduction: - partition = partitionRowwise( - features, phasePlan.hasOperationGraph() ? &phasePlan : nullptr); + partition = partitionRowwise(features, &phasePlan); break; case PhaseBoundaryDomain::IndirectUnderfilledDot: - partition = partitionIndirectDot( - features, phasePlan.hasOperationGraph() ? &phasePlan : nullptr); + partition = partitionIndirectDot(features, &phasePlan); break; } partition.domain = phasePlan.domainName; - if (phasePlan.hasOperationGraph()) { - if (llvm::Error error = - attachCompleteOperationOwnership(partition, phasePlan)) - return std::move(error); - if (anchorPlan) - attachExactAnchorOwnership(partition, *anchorPlan); - deriveStageLiveValues(partition); - if (anchorPlan) - deriveLocalSimtScopeTraffic(partition, *anchorPlan); - } else if (anchorPlan) { - attachAnchorOperationOwnership(partition, *anchorPlan); - } + if (llvm::Error error = + attachCompleteOperationOwnership(partition, phasePlan)) + return std::move(error); + attachExactAnchorOwnership(partition, *anchorPlan); + deriveStageLiveValues(partition); + deriveLocalSimtScopeTraffic(partition, *anchorPlan); return partition; } @@ -1566,125 +1196,78 @@ llvm::Error StageFeatureAnalysis::analyze(StagePartition &partition) const { for (LogicalPhase &phase : partition.phases) { for (LogicalStage &stage : phase.stages) { StageModelFeatures &facts = stage.features; - if (partition.operationOwnershipComplete) { - const double activeLaneRatio = facts.activeLaneRatio; - facts = StageModelFeatures{}; - facts.activeLaneRatio = activeLaneRatio; - llvm::DenseSet owned; - for (Operation *root : stage.operations) - collectOwnedOperationTree(root, owned); - bool hasMemory = false; - int64_t algorithmLoopCount = 0; - for (Operation *operation : owned) { - const llvm::StringRef name = operation->getName().getStringRef(); - if (name == "scf.for" || name == "scf.while") { - facts.hasLoop = true; - ++facts.loopBackedgeCount; - if (!operation->hasAttr("ta.auto_blockify_v1.loop")) - ++algorithmLoopCount; - if (!operation->hasAttr("ta.auto_blockify_v1.loop") && - operation->getNumRegions() > 0 && - !operation->getRegion(0).empty()) { - Block &body = operation->getRegion(0).front(); - const unsigned firstCarriedArgument = name == "scf.for" ? 1 : 0; - for (unsigned argumentIndex = firstCarriedArgument; - argumentIndex < body.getNumArguments(); ++argumentIndex) { - BlockArgument argument = body.getArgument(argumentIndex); - if (argument.use_empty()) - continue; - if (isPointerLikeType(argument.getType()) || - isAddressOnlyLoopValue(argument)) - facts.hasPointerInduction = true; - else - facts.hasLoopCarriedDataDependency = true; - } - // The scf.for induction variable itself commonly feeds pointer - // arithmetic. Record that work without treating it as a data - // recurrence. - if (name == "scf.for" && body.getNumArguments() > 0 && - isAddressOnlyLoopValue(body.getArgument(0))) + const double activeLaneRatio = facts.activeLaneRatio; + facts = StageModelFeatures{}; + facts.activeLaneRatio = activeLaneRatio; + llvm::DenseSet owned; + for (Operation *root : stage.operations) + collectOwnedOperationTree(root, owned); + bool hasMemory = false; + int64_t algorithmLoopCount = 0; + for (Operation *operation : owned) { + const llvm::StringRef name = operation->getName().getStringRef(); + if (name == "scf.for" || name == "scf.while") { + facts.hasLoop = true; + ++facts.loopBackedgeCount; + if (!operation->hasAttr("ta.auto_blockify_v1.loop")) + ++algorithmLoopCount; + if (!operation->hasAttr("ta.auto_blockify_v1.loop") && + operation->getNumRegions() > 0 && + !operation->getRegion(0).empty()) { + Block &body = operation->getRegion(0).front(); + const unsigned firstCarriedArgument = name == "scf.for" ? 1 : 0; + for (unsigned argumentIndex = firstCarriedArgument; + argumentIndex < body.getNumArguments(); ++argumentIndex) { + BlockArgument argument = body.getArgument(argumentIndex); + if (argument.use_empty()) + continue; + if (isPointerLikeType(argument.getType()) || + isAddressOnlyLoopValue(argument)) facts.hasPointerInduction = true; + else + facts.hasLoopCarriedDataDependency = true; } + if (name == "scf.for" && body.getNumArguments() > 0 && + isAddressOnlyLoopValue(body.getArgument(0))) + facts.hasPointerInduction = true; } - if (name == "scf.if" || name == "cf.cond_br") { - ++facts.conditionalBranchCount; - ++facts.divergentBranchCount; - } - if (name.contains("barrier") || name.contains("sync")) - ++facts.synchronizationCount; - if (name == "tt.load" || name == "tt.store" || name == "tt.gather" || - name.starts_with("tt.atomic")) { - hasMemory = true; - facts.hasIndirectMemory |= - isLoadedIndexDependentMemoryOp(operation) || - name == "tt.gather" || name.starts_with("tt.atomic"); - } - facts.hasReduction |= name == "tt.reduce" || name == "tt.scan" || - name == "linalg.reduce"; - facts.hasDot |= name == "tt.dot" || name.contains("matmul") || - name.contains("mmad"); - facts.hasConversionPack |= - name == "arith.extf" || name == "arith.truncf" || - name == "arith.fptosi" || name == "arith.fptoui" || - name == "arith.sitofp" || name == "arith.uitofp" || - name == "tt.fp_to_fp" || name.contains("convert") || - name.contains("pack") || name.contains("unpack"); } - facts.hasContiguousMemory = hasMemory && !facts.hasIndirectMemory; - // Counts consumed by StageCostEvaluator are per logical iteration. - // Multiple sibling recurrence loops are flattened into one Stage - // iteration space, so identical control events must be normalized by - // the number of sibling loops rather than charged once per loop on - // every flattened iteration. - if (algorithmLoopCount > 0 && stage.iterationCount > 1) { - if (facts.hasLoopCarriedDataDependency) - facts.parallelRecurrenceGroupCount = algorithmLoopCount; - facts.loopBackedgeCount = 1; - facts.conditionalBranchCount = std::max( - facts.conditionalBranchCount > 0 ? 1 : 0, - facts.conditionalBranchCount / algorithmLoopCount); - facts.divergentBranchCount = std::max( - facts.divergentBranchCount > 0 ? 1 : 0, - facts.divergentBranchCount / algorithmLoopCount); + if (name == "scf.if" || name == "cf.cond_br") { + ++facts.conditionalBranchCount; + ++facts.divergentBranchCount; } - facts.source = - "exact post-layout/post-AutoBlockify-V1 TTIR operation graph"; - if (!facts.isValid()) - return llvm::createStringError(std::errc::invalid_argument, - "Stage '%s' has invalid features", - stage.id.c_str()); - continue; + if (name.contains("barrier") || name.contains("sync")) + ++facts.synchronizationCount; + if (name == "tt.load" || name == "tt.store" || name == "tt.gather" || + name.starts_with("tt.atomic")) { + hasMemory = true; + facts.hasIndirectMemory |= + isLoadedIndexDependentMemoryOp(operation) || + name == "tt.gather" || name.starts_with("tt.atomic"); + } + facts.hasReduction |= + name == "tt.reduce" || name == "tt.scan" || name == "linalg.reduce"; + facts.hasDot |= name == "tt.dot" || name.contains("matmul") || + name.contains("mmad"); + facts.hasConversionPack |= + name == "arith.extf" || name == "arith.truncf" || + name == "arith.fptosi" || name == "arith.fptoui" || + name == "arith.sitofp" || name == "arith.uitofp" || + name == "tt.fp_to_fp" || name.contains("convert") || + name.contains("pack") || name.contains("unpack"); + } + facts.hasContiguousMemory = hasMemory && !facts.hasIndirectMemory; + if (algorithmLoopCount > 0 && stage.iterationCount > 1) { + if (facts.hasLoopCarriedDataDependency) + facts.parallelRecurrenceGroupCount = algorithmLoopCount; + facts.loopBackedgeCount = 1; + facts.conditionalBranchCount = std::max( + facts.conditionalBranchCount > 0 ? 1 : 0, + facts.conditionalBranchCount / algorithmLoopCount); + facts.divergentBranchCount = + std::max(facts.divergentBranchCount > 0 ? 1 : 0, + facts.divergentBranchCount / algorithmLoopCount); } - facts.hasLoop = - stage.iterationCount > 1 || - stage.costModelKind == StageCostModelKind::AutoBlockifyLoop || - stage.costModelKind == StageCostModelKind::IndependentPipelinedLoop || - stage.costModelKind == StageCostModelKind::LoopCarriedRecurrence; - facts.loopBackedgeCount = facts.hasLoop ? 1 : 0; - facts.hasLoopCarriedDataDependency = - stage.costModelKind == StageCostModelKind::LoopCarriedRecurrence; - facts.hasPointerInduction = - facts.hasLoop && !facts.hasLoopCarriedDataDependency; - facts.hasContiguousMemory = - llvm::is_contained({StageCostModelKind::ContinuousTileMemory, - StageCostModelKind::ContinuousTileStore, - StageCostModelKind::ContinuousShortLoad, - StageCostModelKind::CachePolicyStore}, - stage.costModelKind); - facts.hasIndirectMemory = - llvm::is_contained({StageCostModelKind::IndirectScalarMemory, - StageCostModelKind::IndirectGatherMemory}, - stage.costModelKind); - facts.hasReduction = - llvm::is_contained({StageCostModelKind::LoopCarriedRecurrence, - StageCostModelKind::RowwiseReduction}, - stage.costModelKind); - facts.hasDot = llvm::is_contained({StageCostModelKind::CubeRoofline, - StageCostModelKind::TinyCubeRoofline}, - stage.costModelKind); - facts.hasConversionPack = - stage.costModelKind == StageCostModelKind::ConversionPack; - facts.source = "feature-summary fallback structural facts"; if (!facts.isValid()) return llvm::createStringError(std::errc::invalid_argument, "Stage '%s' has invalid features", @@ -1698,67 +1281,43 @@ llvm::Error StageKindClassifier::analyze(StagePartition &partition, int64_t tinyDotFlopsMax) const { if (!partition.operationOwnershipComplete) return llvm::Error::success(); + auto compatible = [](StageCostModelKind kind, + const StageModelFeatures &facts) { + switch (kind) { + case StageCostModelKind::LoopCarriedRecurrence: + return facts.hasLoopCarriedDataDependency; + case StageCostModelKind::IndependentPipelinedLoop: + return facts.hasLoop && !facts.hasLoopCarriedDataDependency; + case StageCostModelKind::RowwiseReduction: + return facts.hasReduction; + case StageCostModelKind::CubeRoofline: + case StageCostModelKind::TinyCubeRoofline: + return facts.hasDot; + case StageCostModelKind::IndirectScalarMemory: + case StageCostModelKind::IndirectGatherMemory: + return facts.hasIndirectMemory; + case StageCostModelKind::ContinuousTileMemory: + case StageCostModelKind::ContinuousTileStore: + case StageCostModelKind::ContinuousShortLoad: + case StageCostModelKind::CachePolicyStore: + return facts.hasContiguousMemory; + case StageCostModelKind::ConversionPack: + return facts.hasConversionPack; + default: + return true; + } + }; for (LogicalPhase &phase : partition.phases) { for (LogicalStage &stage : phase.stages) { const StageModelFeatures &facts = stage.features; - // Auxiliary scalar/index/predicate and memory work may live in a - // specialized Stage, but two independent dominant formulas may not. - // Reaching this check means StageBoundaryAnalysis missed a structural - // cut. Silently selecting the first matching Kind would hide or - // double-count work, so report a stable boundary diagnostic instead. - // Conversion operations are not a split condition by themselves: - // predicate-to-float and accumulator casts are often auxiliary work of - // a recurrence/reduction/dot Stage. ConversionPack becomes dominant - // only when no stronger structure owns the Stage (see deriveKind()). - const bool requiresSplit = - facts.hasDot && (facts.hasReduction || facts.hasIndirectMemory || - facts.hasLoopCarriedDataDependency); - if (requiresSplit) + if (facts.hasDot && (facts.hasReduction || facts.hasIndirectMemory || + facts.hasLoopCarriedDataDependency)) return llvm::createStringError( std::errc::invalid_argument, - "requires_split: Stage '%s' owns incompatible dominant " - "structures (carried=%d, indirect=%d, reduction=%d, dot=%d, " - "conversion=%d, roots=%zu)", - stage.id.c_str(), facts.hasLoopCarriedDataDependency, - facts.hasIndirectMemory, facts.hasReduction, facts.hasDot, - facts.hasConversionPack, stage.operations.size()); - auto kindMatchesFacts = [&](StageCostModelKind kind) { - switch (kind) { - case StageCostModelKind::AutoBlockifyDispatch: - case StageCostModelKind::AutoBlockifyLoop: - case StageCostModelKind::ScalarIssue: - case StageCostModelKind::ScalarControl: - case StageCostModelKind::ScalarMath: - case StageCostModelKind::IndexGeneration: - case StageCostModelKind::PredicateMask: - case StageCostModelKind::LoopPredicate: - return true; - case StageCostModelKind::LoopCarriedRecurrence: - return facts.hasLoopCarriedDataDependency; - case StageCostModelKind::IndependentPipelinedLoop: - return facts.hasLoop && !facts.hasLoopCarriedDataDependency; - case StageCostModelKind::RowwiseReduction: - return facts.hasReduction; - case StageCostModelKind::CubeRoofline: - case StageCostModelKind::TinyCubeRoofline: - return facts.hasDot; - case StageCostModelKind::IndirectScalarMemory: - case StageCostModelKind::IndirectGatherMemory: - return facts.hasIndirectMemory; - case StageCostModelKind::ContinuousTileMemory: - case StageCostModelKind::ContinuousTileStore: - case StageCostModelKind::ContinuousShortLoad: - case StageCostModelKind::CachePolicyStore: - return facts.hasContiguousMemory; - case StageCostModelKind::ConversionPack: - return facts.hasConversionPack; - } - return false; - }; - auto deriveKind = [&]() { - // Preserve semantic specializations when their defining evidence is - // present. Otherwise classify from the exact Stage operation graph; - // no workload name or experiment identity participates here. + "requires_split: Stage '%s' owns incompatible dominant structures", + stage.id.c_str()); + + auto derive = [&]() { if (facts.hasDot) return stage.workload.dotFlops * stage.iterationCount <= static_cast( @@ -1782,106 +1341,26 @@ llvm::Error StageKindClassifier::analyze(StagePartition &partition, : StageCostModelKind::ContinuousTileMemory; return StageCostModelKind::ScalarIssue; }; - if (!kindMatchesFacts(stage.costModelKind)) { - stage.costModelKind = deriveKind(); - if (stage.costModelKind == StageCostModelKind::IndependentPipelinedLoop) - stage.scheduleKind = StageScheduleKind::IndependentPipelined; - else if (stage.costModelKind == - StageCostModelKind::LoopCarriedRecurrence) - stage.scheduleKind = StageScheduleKind::LoopCarriedSerial; - } - const StageCostModelKind kind = stage.costModelKind; - auto mismatch = [&]() { + if (!compatible(stage.costModelKind, facts)) + stage.costModelKind = derive(); + if (!compatible(stage.costModelKind, facts) || + (stage.costModelKind == StageCostModelKind::TinyCubeRoofline && + stage.workload.dotFlops * stage.iterationCount > + static_cast(std::max(1, tinyDotFlopsMax)))) return llvm::createStringError( std::errc::invalid_argument, - "Stage '%s' operation graph does not match StageCostModelKind " - "'%s' (loop=%d, carried=%d, contiguous=%d, indirect=%d, " - "reduction=%d, dot=%d, conversion=%d, roots=%zu)", - stage.id.c_str(), stringifyStageCostModel(kind).str().c_str(), - facts.hasLoop, facts.hasLoopCarriedDataDependency, - facts.hasContiguousMemory, facts.hasIndirectMemory, - facts.hasReduction, facts.hasDot, facts.hasConversionPack, - stage.operations.size()); - }; - switch (kind) { - case StageCostModelKind::AutoBlockifyDispatch: - case StageCostModelKind::AutoBlockifyLoop: - break; - case StageCostModelKind::LoopCarriedRecurrence: - if (!facts.hasLoopCarriedDataDependency) - return mismatch(); - break; - case StageCostModelKind::IndependentPipelinedLoop: - if (!facts.hasLoop || facts.hasLoopCarriedDataDependency) - return mismatch(); - break; - case StageCostModelKind::RowwiseReduction: - if (!facts.hasReduction) - return mismatch(); - break; - case StageCostModelKind::CubeRoofline: - if (!facts.hasDot) - return mismatch(); - break; - case StageCostModelKind::TinyCubeRoofline: - if (!facts.hasDot || - stage.workload.dotFlops * stage.iterationCount > - static_cast(std::max(1, tinyDotFlopsMax))) - return mismatch(); - break; - case StageCostModelKind::IndirectScalarMemory: - case StageCostModelKind::IndirectGatherMemory: - if (!facts.hasIndirectMemory) - return mismatch(); - break; - case StageCostModelKind::ContinuousTileMemory: - case StageCostModelKind::ContinuousTileStore: - case StageCostModelKind::ContinuousShortLoad: - case StageCostModelKind::CachePolicyStore: - if (!facts.hasContiguousMemory) - return mismatch(); - break; - case StageCostModelKind::ConversionPack: - if (!facts.hasConversionPack) - return mismatch(); - break; - case StageCostModelKind::ScalarIssue: - case StageCostModelKind::ScalarControl: - case StageCostModelKind::ScalarMath: - case StageCostModelKind::IndexGeneration: - case StageCostModelKind::PredicateMask: - case StageCostModelKind::LoopPredicate: - // These kinds may legitimately contain auxiliary address, predicate, - // control, or short metadata memory work. Their dominant semantics - // is established by the contiguous boundary plan. - break; - } + "Stage '%s' operation graph does not match StageCostModelKind '%s'", + stage.id.c_str(), + stringifyStageCostModel(stage.costModelKind).str().c_str()); + if (stage.costModelKind == StageCostModelKind::IndependentPipelinedLoop) + stage.scheduleKind = StageScheduleKind::IndependentPipelined; + else if (stage.costModelKind == StageCostModelKind::LoopCarriedRecurrence) + stage.scheduleKind = StageScheduleKind::LoopCarriedSerial; } } return llvm::Error::success(); } -llvm::Error -StageWorkloadAnalysis::verify(const StagePartition &partition, - const StageWorkload &kernelWorkload) const { - if (!kernelWorkload.isFiniteAndNonNegative()) - return llvm::createStringError(std::errc::invalid_argument, - "kernel StageWorkload is invalid"); - const StageWorkload owned = accumulatePartition(partition); - if (!near(owned.scalarOperations, kernelWorkload.scalarOperations) || - !near(owned.loadBytes, kernelWorkload.loadBytes) || - !near(owned.storeBytes, kernelWorkload.storeBytes) || - !near(owned.predicateElements, kernelWorkload.predicateElements) || - !near(owned.shuffleLaneSteps, kernelWorkload.shuffleLaneSteps) || - !near(owned.dotFlops, kernelWorkload.dotFlops) || - !near(totalOperationElements(owned), - totalOperationElements(kernelWorkload))) - return llvm::createStringError( - std::errc::invalid_argument, - "StagePartition does not conserve post-transform TTIR workload"); - return llvm::Error::success(); -} - llvm::Error StageWorkloadAnalysis::analyze(StagePartition &partition) const { if (!partition.operationOwnershipComplete) return llvm::createStringError( @@ -1911,8 +1390,7 @@ llvm::Error StageWorkloadAnalysis::analyze(StagePartition &partition) const { } llvm::Error -StagePartitionVerifier::verify(const StagePartition &partition, - const StageWorkload &kernelWorkload) const { +StagePartitionVerifier::verify(const StagePartition &partition) const { if (partition.phases.empty()) return llvm::createStringError(std::errc::invalid_argument, "StagePartition has no Phase"); @@ -1969,9 +1447,11 @@ StagePartitionVerifier::verify(const StagePartition &partition, return llvm::createStringError( std::errc::invalid_argument, "StagePartition operation ownership is incomplete"); - if (partition.operationOwnershipComplete) - return llvm::Error::success(); - return StageWorkloadAnalysis().verify(partition, kernelWorkload); + if (!partition.operationOwnershipComplete) + return llvm::createStringError( + std::errc::invalid_argument, + "StagePartition requires complete TTIR operation ownership"); + return llvm::Error::success(); } llvm::Error @@ -1979,6 +1459,11 @@ StageModeLegalityAnalysis::analyze(StagePartition &partition, int64_t maximumSuperblockFactor, bool scopeSuperblockMaterializable) const { const int64_t maximum = std::clamp(maximumSuperblockFactor, 1, 4); + // Local and whole-kernel SuperBlock candidates consume the same SIMT warp + // resources. Do not regenerate F4 here after evaluateStageModel has + // already reduced the target maximum to F2 for num_warps=32 (or to F1 for + // a smaller runtime grid). + const int64_t localMaximum = scopeSuperblockMaterializable ? maximum : 1; for (LogicalPhase &phase : partition.phases) { for (LogicalStage &stage : phase.stages) { stage.simdLegal = true; @@ -1994,18 +1479,21 @@ StageModeLegalityAnalysis::analyze(StagePartition &partition, if (maximum >= 4) stage.legalSimtFactors.push_back(4); if (stage.localSimtMaterializable) { - // A mixed F2/F4 implementation is legal only when backend integration - // can wrap the complete logical-program body with AutoBlockify V1. - // The local scope itself remains a single-mode Stage; factor changes - // its implementation, not the Stage boundary. - stage.localSimtFactors = scopeSuperblockMaterializable - ? stage.legalSimtFactors - : std::vector{1}; + // The ABI-v2 scope materializer batches complete logical programs + // around this Stage. F2/F4 therefore does not require multiple + // recurrence groups inside one logical program; that older W2/W4 + // interpretation was only warp widening, not a SuperBlock. + stage.localSimtFactors = {1}; + if (scopeSuperblockMaterializable) + for (int64_t factor : {2, 4}) + if (factor <= localMaximum) + stage.localSimtFactors.push_back(factor); } if (stage.localSimtMaterializable && (stage.localSimtFactors.empty() || llvm::any_of(stage.localSimtFactors, [&](int64_t factor) { - return !llvm::is_contained(stage.legalSimtFactors, factor); + return factor < 1 || factor > localMaximum || + (factor != 1 && factor != 2 && factor != 4); }))) return llvm::createStringError( std::errc::invalid_argument, @@ -2015,35 +1503,6 @@ StageModeLegalityAnalysis::analyze(StagePartition &partition, return llvm::Error::success(); } -llvm::Expected> -StagePartitioner::partition(const SimdSimtFeatureSummary &features, - const StagePartitionerOptions &options) const { - auto phasePlan = PhaseBoundaryAnalysis().analyze(features, options); - if (!phasePlan) - return phasePlan.takeError(); - if (!*phasePlan) - return std::optional{}; - auto result = StageBoundaryAnalysis().analyze(**phasePlan, features); - if (!result) - return result.takeError(); - - StageFeatureAnalysis featureAnalysis; - if (llvm::Error error = featureAnalysis.analyze(*result)) - return std::move(error); - if (llvm::Error error = - StageKindClassifier().analyze(*result, options.tinyDotFlopsMax)) - return std::move(error); - StageModeLegalityAnalysis legalityAnalysis; - if (llvm::Error error = - legalityAnalysis.analyze(*result, options.maximumSuperblockFactor, - options.scopeSuperblockMaterializable)) - return std::move(error); - if (llvm::Error error = StagePartitionVerifier().verify( - *result, buildKernelStageWorkload(features))) - return std::move(error); - return std::optional{std::move(*result)}; -} - llvm::Expected> StagePartitioner::partition(ModuleOp module, const SimtAnchorPlan &anchorPlan, const SimdSimtFeatureSummary &features, @@ -2072,38 +1531,7 @@ StagePartitioner::partition(ModuleOp module, const SimtAnchorPlan &anchorPlan, legalityAnalysis.analyze(*result, options.maximumSuperblockFactor, options.scopeSuperblockMaterializable)) return std::move(error); - if (llvm::Error error = StagePartitionVerifier().verify( - *result, buildKernelStageWorkload(features))) - return std::move(error); - return std::optional{std::move(*result)}; -} - -llvm::Expected> -StagePartitioner::partition(const SimdSimtFeatureSummary &features, - const StagePartitionerOptions &options, - const SimtAnchorPlan &anchorPlan) const { - auto phasePlan = PhaseBoundaryAnalysis().analyze(features, options); - if (!phasePlan) - return phasePlan.takeError(); - if (!*phasePlan) - return std::optional{}; - auto result = - StageBoundaryAnalysis().analyze(**phasePlan, features, &anchorPlan); - if (!result) - return result.takeError(); - StageFeatureAnalysis featureAnalysis; - if (llvm::Error error = featureAnalysis.analyze(*result)) - return std::move(error); - if (llvm::Error error = - StageKindClassifier().analyze(*result, options.tinyDotFlopsMax)) - return std::move(error); - StageModeLegalityAnalysis legalityAnalysis; - if (llvm::Error error = - legalityAnalysis.analyze(*result, options.maximumSuperblockFactor, - options.scopeSuperblockMaterializable)) - return std::move(error); - if (llvm::Error error = StagePartitionVerifier().verify( - *result, buildKernelStageWorkload(features))) + if (llvm::Error error = StagePartitionVerifier().verify(*result)) return std::move(error); return std::optional{std::move(*result)}; } diff --git a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/CMakeLists.txt b/third_party/ascend/costmodel/lib/AscendModel/RouteModel/CMakeLists.txt index 6e73a954cf..bda18f3d26 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/CMakeLists.txt +++ b/third_party/ascend/costmodel/lib/AscendModel/RouteModel/CMakeLists.txt @@ -5,13 +5,13 @@ add_compile_definitions( ) add_mlir_library(AscendModelRouteModel - SimtAnchorAnalysis.cpp - StagePartitioner.cpp + ../Analysis/SimtAnchorAnalysis.cpp + ../Analysis/StagePartitioner.cpp StageCostModels.cpp StageRouteCostModel.cpp SimdSimtCostModel.cpp - Transforms/SelectSimdSimtCostModel.cpp - Transforms/MaterializeSimtScopes.cpp + ../Transforms/SelectSimdSimtCostModel.cpp + ../Transforms/MaterializeSimtScopes.cpp ADDITIONAL_HEADER_DIRS ${TRITON_ASCEND_COSTMODEL_SOURCE_DIR}/include/AscendModel/RouteModel diff --git a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/SimdSimtCostModel.cpp b/third_party/ascend/costmodel/lib/AscendModel/RouteModel/SimdSimtCostModel.cpp index a4a443cc28..9e681d8240 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/SimdSimtCostModel.cpp +++ b/third_party/ascend/costmodel/lib/AscendModel/RouteModel/SimdSimtCostModel.cpp @@ -7,10 +7,10 @@ //===----------------------------------------------------------------------===// #include "AscendModel/RouteModel/SimdSimtCostModel.h" +#include "AscendModel/Analysis/SimtAnchorAnalysis.h" +#include "AscendModel/Analysis/StagePartitioner.h" #include "AscendModel/Profile/MicrobenchmarkProfile.h" -#include "AscendModel/RouteModel/SimtAnchorAnalysis.h" #include "AscendModel/RouteModel/StageCostModels.h" -#include "AscendModel/RouteModel/StagePartitioner.h" #include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/BuiltinTypes.h" @@ -48,106 +48,19 @@ constexpr llvm::StringLiteral kAllSimd = "all_simd"; constexpr llvm::StringLiteral kAllSimtOnly = "all_simt_only"; constexpr llvm::StringLiteral kMixedSimdSimt = "mixed_simd_simt"; -struct OpProfile { - double throughput = 0.0; - double factor = 1.0; - std::string confidence = "none"; -}; - struct StructuralProfile { - double irregularPerDensity = 0.0; - double irregularCap = 0.0; - double tinyDotIrregularPerDensity = 0.0; - double tinyDotIrregularCap = 0.0; - double perMaskRank = 0.0; - double maskCap = 0.0; - double perWeightedReduction = 0.0; - double reductionCap = 0.0; - double perStaticLoopTrip = 0.0; - double loopCap = 0.0; - double controlFlow = 0.0; - double rank1IndirectVectorReduction = 0.0; - double tinyDot = 0.0; int64_t tinyDotFlopsMax = 0; }; -struct MixedSetupFallbackProfile { - int64_t numWarps = 0; - double emptySimtSetupCycles = 0.0; -}; - -struct StageResourceProfile { - double scalarOperationsPerCycle = 0.0; - double issueOperationsPerCycle = 0.0; - double spillTransactionsPerCycle = 0.0; - double indirectLoadTransactionsPerCycle = 0.0; - double indirectStoreTransactionsPerCycle = 0.0; - double indirectDependencyLatencyCycles = 0.0; - StageControlFlowRates controlFlow; - - bool isValid() const { - const std::array positive = { - scalarOperationsPerCycle, issueOperationsPerCycle, - spillTransactionsPerCycle, indirectLoadTransactionsPerCycle, - indirectStoreTransactionsPerCycle}; - return llvm::all_of(positive, - [](double value) { - return std::isfinite(value) && value > 0.0; - }) && - std::isfinite(indirectDependencyLatencyCycles) && - indirectDependencyLatencyCycles >= 0.0 && - controlFlow.isFiniteAndNonNegative(); - } -}; - struct CandidateProfile { - std::string profileVersion; - std::string target; - std::vector compatibleTargets; + HardwareProfile hardware; std::string scoreUnit; std::string contentSha256; std::string selectionContentSha256; std::string microbenchmarkProfileVersion; std::string microbenchmarkProfileTarget; std::string microbenchmarkContentSha256; - - double programIssueScale = 1.0; StructuralProfile structural; - int64_t simdVectorWidthBits = 2048; - double simdSetupCycles = 0.0; - llvm::StringMap simdOps; - double simdMte2BytesPerCycle = 0.0; - double simdMte3BytesPerCycle = 0.0; - std::string simdMemoryConfidence = "none"; - double simdDotSetupCycles = 0.0; - double simdDotFlopsPerCycle = 0.0; - std::string simdDotConfidence = "none"; - StageResourceProfile simdStageResources; - - int64_t simtWarpSize = 32; - double simtSetupCycles = 0.0; - std::string simtSetupConfidence = "none"; - llvm::StringMap simtOps; - double simtDotSetupCycles = 0.0; - double simtDotFlopsPerCycle = 0.0; - std::string simtDotConfidence = "none"; - double simtPredicateRate = 0.0; - double simtShuffleRate = 0.0; - std::string simtShuffleConfidence = "none"; - double simtLoadWarpRate = 0.0; - double simtStoreWarpRate = 0.0; - std::string simtMemoryConfidence = "none"; - StageResourceProfile simtStageResources; - int64_t superblockUsefulFactorLimit = 1; - int64_t superblockPersistentStatePressureFreeFactor = 1; - double superblockPersistentStateBytesPerCycle = 0.0; - double scopeHandoffFixedDirectionalCycles = 0.0; - double scopeSimdUbLoadBytesPerCycle = 0.0; - double scopeSimdUbStoreBytesPerCycle = 0.0; - double scopeSimtUbLoadBytesPerThreadPerCycle = 0.0; - double scopeSimtUbStoreBytesPerThreadPerCycle = 0.0; - std::vector mixedSetupFallbacks; - std::string mixedSetupFallbackConfidence = "none"; }; /// Small fail-fast facade around llvm::json. It permits a readable profile @@ -165,16 +78,6 @@ class ProfileJSONReader { return nullptr; } - const llvm::json::Array *array(const llvm::json::Object &parent, - llvm::StringRef key, llvm::StringRef context) { - if (failed()) - return nullptr; - if (const auto *value = parent.getArray(key)) - return value; - setError(context + "." + key + " must be an array"); - return nullptr; - } - double number(const llvm::json::Object &parent, llvm::StringRef key, llvm::StringRef context) { if (failed()) @@ -205,14 +108,6 @@ class ProfileJSONReader { return {}; } - std::string optionalString(const llvm::json::Object &parent, - llvm::StringRef key, - llvm::StringRef defaultValue = {}) { - if (auto value = parent.getString(key)) - return value->str(); - return defaultValue.str(); - } - double optionalNumber(const llvm::json::Object &parent, llvm::StringRef key, double defaultValue) { if (auto value = parent.getNumber(key)) @@ -220,13 +115,6 @@ class ProfileJSONReader { return defaultValue; } - int64_t optionalInteger(const llvm::json::Object &parent, llvm::StringRef key, - int64_t defaultValue) { - if (auto value = parent.getInteger(key)) - return *value; - return defaultValue; - } - bool failed() const { return !error.empty(); } llvm::StringRef getError() const { return error; } @@ -239,155 +127,6 @@ class ProfileJSONReader { std::string error; }; -static llvm::json::Object toJSON(const llvm::StringMap &values) { - llvm::json::Object result; - for (const auto &entry : values) - result[entry.first()] = entry.second; - return result; -} - -static llvm::json::Object toJSON(const llvm::StringMap &values) { - llvm::json::Object result; - for (const auto &entry : values) - result[entry.first()] = entry.second; - return result; -} - -template -static void initializeWorkMaps(SummaryT &features) { - for (llvm::StringRef key : - {"load", "store", "reduce", "scan", "gather", "histogram", "atomic", - "add", "sub", "mul", "div", "max", "abs", "exp", "log", "cmp", "select", - "cast", "clamp"}) { - features.weightedOps[key] = 0; - features.opElements[key] = 0; - } -} - -static std::string typeToString(Type type) { - std::string text; - llvm::raw_string_ostream os(text); - os << type; - os.flush(); - return text; -} - -static bool isPointerType(Type type) { - if (auto tensor = dyn_cast(type)) - type = tensor.getElementType(); - return llvm::StringRef(typeToString(type)).contains("!tt.ptr"); -} - -/// Return true only when a loop-carried value is used exclusively to derive -/// load/store addresses (plus the loop yield). This recognizes pointer and -/// integer-offset induction without confusing it with a recurrence whose -/// value feeds arithmetic, predicates, or stored data. -static bool isAddressOnlyLoopCarriedValue(Value root) { - llvm::SmallVector worklist{root}; - llvm::DenseSet visited; - bool reachesAddressUse = false; - while (!worklist.empty()) { - Value value = worklist.pop_back_val(); - if (!visited.insert(value).second) - continue; - for (OpOperand &use : value.getUses()) { - Operation *user = use.getOwner(); - llvm::StringRef name = user->getName().getStringRef(); - if (name == "scf.yield") - continue; - if ((name == "tt.load" || name == "tt.store") && - use.getOperandNumber() == 0) { - reachesAddressUse = true; - continue; - } - const bool addressForwarding = - name == "tt.addptr" || name == "tt.splat" || name == "tt.broadcast" || - name == "tt.expand_dims" || name == "arith.addi" || - name == "arith.subi" || name == "arith.muli" || - name == "arith.index_cast"; - if (!addressForwarding) - return false; - if (name == "tt.addptr") - reachesAddressUse = true; - for (Value result : user->getResults()) - worklist.push_back(result); - } - } - return reachesAddressUse; -} - -static Type getElementType(Type type) { - if (auto tensor = dyn_cast(type)) - return tensor.getElementType(); - return type; -} - -static int64_t parseTypeBitWidth(llvm::StringRef text) { - for (size_t index = 0; index + 1 < text.size(); ++index) { - if (text[index] != 'f' && text[index] != 'i') - continue; - size_t end = index + 1; - while (end < text.size() && llvm::isDigit(text[end])) - ++end; - if (end == index + 1) - continue; - int64_t width = 0; - if (!text.slice(index + 1, end).getAsInteger(10, width) && width > 0) - return width; - } - return 0; -} - -static int64_t getTypeBitWidth(Type type, int64_t defaultWidth = 32) { - type = getElementType(type); - if (auto integer = dyn_cast(type)) - return integer.getWidth(); - if (auto floating = dyn_cast(type)) - return floating.getWidth(); - int64_t parsed = parseTypeBitWidth(typeToString(type)); - return parsed > 0 ? parsed : defaultWidth; -} - -static bool isMaskTensorType(Type type) { - auto tensor = dyn_cast(type); - if (!tensor) - return false; - auto integer = dyn_cast(tensor.getElementType()); - return integer && integer.getWidth() == 1; -} - -static int64_t getStaticNumElements(Type type) { - auto tensor = dyn_cast(type); - if (!tensor) - return 1; - int64_t count = 1; - for (int64_t dim : tensor.getShape()) { - if (ShapedType::isDynamic(dim) || dim <= 0) - return 1; - if (count > std::numeric_limits::max() / dim) - return std::numeric_limits::max(); - count *= dim; - } - return std::max(1, count); -} - -static double getStaticTensorBytes(Type type) { - auto tensor = dyn_cast(type); - if (!tensor || !tensor.hasStaticShape()) - return 0.0; - return static_cast(getStaticNumElements(type)) * - getTypeBitWidth(tensor.getElementType()) / 8.0; -} - -static int64_t getOperationElements(Operation *op) { - int64_t elements = 1; - for (Type type : op->getOperandTypes()) - elements = std::max(elements, getStaticNumElements(type)); - for (Type type : op->getResultTypes()) - elements = std::max(elements, getStaticNumElements(type)); - return elements; -} - static std::optional getConstantInteger(Value value) { Operation *definingOp = value.getDefiningOp(); if (!definingOp || definingOp->getName().getStringRef() != "arith.constant") @@ -452,136 +191,11 @@ static int64_t getLoopMultiplier( return multiplier; } -static bool isCastOp(llvm::StringRef name) { - return name.starts_with("arith.ext") || name.starts_with("arith.trunc") || - name == "arith.sitofp" || name == "arith.uitofp" || - name == "arith.fptosi" || name == "arith.fptoui" || - name.starts_with("arith.index_cast"); -} - -static llvm::StringRef classifyWeightedOp(llvm::StringRef name) { - if (name == "tt.load") - return "load"; - if (name == "tt.store") - return "store"; - if (name == "tt.reduce") - return "reduce"; - if (name == "tt.scan" || name == "tt.associative_scan") - return "scan"; - if (name == "tt.gather") - return "gather"; - if (name == "tt.histogram") - return "histogram"; - if (name.starts_with("tt.atomic")) - return "atomic"; - if (name == "arith.addf" || name == "arith.addi") - return "add"; - if (name == "arith.subf" || name == "arith.subi") - return "sub"; - if (name == "arith.mulf" || name == "arith.muli") - return "mul"; - if (name == "arith.divf" || name == "arith.divsi" || name == "arith.divui") - return "div"; - if (name == "arith.maxnumf" || name == "arith.maxf" || - name == "arith.maxsi" || name == "arith.maxui") - return "max"; - if (name == "math.absf" || name == "math.absi") - return "abs"; - if (name == "math.exp") - return "exp"; - if (name == "math.log") - return "log"; - if (name == "arith.cmpf" || name == "arith.cmpi") - return "cmp"; - if (name == "arith.select") - return "select"; - if (isCastOp(name)) - return "cast"; - if (name.starts_with("tt.clamp")) - return "clamp"; - return {}; -} - -static void appendUnique(std::vector &values, - llvm::StringRef value) { - if (llvm::find(values, value.str()) == values.end()) - values.push_back(value.str()); -} - -static int confidenceRank(llvm::StringRef confidence) { - if (confidence == "high") - return 3; - if (confidence == "medium") - return 2; - if (confidence == "low") - return 1; - return 0; -} - -static std::string minimumConfidence(llvm::ArrayRef values) { - if (values.empty()) - return "none"; - return *std::min_element(values.begin(), values.end(), - [](const std::string &lhs, const std::string &rhs) { - return confidenceRank(lhs) < confidenceRank(rhs); - }); -} - -static bool wildcardMatch(llvm::StringRef pattern, llvm::StringRef value) { - size_t patternIndex = 0; - size_t valueIndex = 0; - size_t starIndex = llvm::StringRef::npos; - size_t retryValueIndex = 0; - while (valueIndex < value.size()) { - if (patternIndex < pattern.size() && - (pattern[patternIndex] == '?' || - pattern[patternIndex] == value[valueIndex])) { - ++patternIndex; - ++valueIndex; - continue; - } - if (patternIndex < pattern.size() && pattern[patternIndex] == '*') { - starIndex = patternIndex++; - retryValueIndex = valueIndex; - continue; - } - if (starIndex != llvm::StringRef::npos) { - patternIndex = starIndex + 1; - valueIndex = ++retryValueIndex; - continue; - } - return false; - } - while (patternIndex < pattern.size() && pattern[patternIndex] == '*') - ++patternIndex; - return patternIndex == pattern.size(); -} - -static bool targetMatches(const CandidateProfile &profile, - llvm::StringRef actualTarget) { - if (actualTarget.trim().empty()) - return true; - std::string actual = actualTarget.trim().lower(); - std::vector patterns = profile.compatibleTargets; - patterns.push_back(profile.target); - for (std::string pattern : patterns) { - std::replace(pattern.begin(), pattern.end(), ':', '/'); - llvm::SmallVector aliases; - llvm::StringRef(pattern).split(aliases, '/', -1, false); - for (llvm::StringRef alias : aliases) { - std::string lower = alias.trim().lower(); - if (!lower.empty() && wildcardMatch(lower, actual)) - return true; - } - } - return false; -} - static double resolveNumberOrMeasurement( const llvm::json::Object &object, llvm::StringRef numberKey, llvm::StringRef measurementKey, llvm::StringRef expectedUnit, const MicrobenchmarkProfile *microbench, ProfileJSONReader &reader, - llvm::StringRef context, std::string *measurementConfidence = nullptr) { + llvm::StringRef context) { if (auto reference = object.getString(measurementKey)) { if (!microbench) { reader.setError(context + "." + measurementKey + @@ -598,23 +212,17 @@ static double resolveNumberOrMeasurement( reader.setError(llvm::toString(value.takeError())); return 0.0; } - if (measurementConfidence) { - const MicrobenchmarkMeasurement *measurement = - microbench->getMeasurement(*reference); - *measurementConfidence = measurement ? measurement->confidence : "none"; - } return *value; } return reader.number(object, numberKey, context); } -static OpProfile resolveOpProfile(const llvm::json::Object &ops, - llvm::StringRef opName, - llvm::StringRef throughputKey, - llvm::StringRef expectedUnit, - const MicrobenchmarkProfile *microbench, - ProfileJSONReader &reader) { - OpProfile result; +static StageOperationRate +resolveOpProfile(const llvm::json::Object &ops, llvm::StringRef opName, + llvm::StringRef throughputKey, llvm::StringRef expectedUnit, + const MicrobenchmarkProfile *microbench, + ProfileJSONReader &reader) { + StageOperationRate result; const llvm::json::Value *raw = ops.get(opName); if (!raw) { reader.setError("missing operation profile " + opName); @@ -626,20 +234,16 @@ static OpProfile resolveOpProfile(const llvm::json::Object &ops, return result; } if (auto relative = op->getString("relative_to")) { - OpProfile base = resolveOpProfile(ops, *relative, throughputKey, - expectedUnit, microbench, reader); + StageOperationRate base = resolveOpProfile( + ops, *relative, throughputKey, expectedUnit, microbench, reader); result.throughput = base.throughput; result.factor = reader.optionalNumber(*op, "factor", 1.0); - result.confidence = reader.optionalString(*op, "confidence", "low"); return result; } - std::string measuredConfidence = "none"; - result.throughput = resolveNumberOrMeasurement( - *op, throughputKey, "throughput_measurement", expectedUnit, microbench, - reader, opName, &measuredConfidence); + result.throughput = + resolveNumberOrMeasurement(*op, throughputKey, "throughput_measurement", + expectedUnit, microbench, reader, opName); result.factor = reader.optionalNumber(*op, "factor", 1.0); - result.confidence = - reader.optionalString(*op, "confidence", measuredConfidence); return result; } @@ -692,6 +296,43 @@ static std::string resolveProfileReference(llvm::StringRef ownerPath, return resolved.str().str(); } +static void readStageResources(ProfileJSONReader &reader, + const llvm::json::Object &mode, + llvm::StringRef context, + StageModeProfile &profile) { + const auto *resources = reader.object(mode, "stage_resources", context); + if (!resources) + return; + const std::string prefix = (context + ".stage_resources").str(); + profile.scalarOperationsPerCycle = + reader.number(*resources, "scalar_operations_per_system_cycle", prefix); + profile.issueOperationsPerCycle = + reader.number(*resources, "issue_instructions_per_system_cycle", prefix); + profile.spillTransactionsPerCycle = + reader.number(*resources, "spill_transactions_per_system_cycle", prefix); + if (const auto *indirect = + reader.object(*resources, "indirect_memory", prefix)) { + const std::string path = prefix + ".indirect_memory"; + profile.indirectLoadTransactionsPerCycle = + reader.number(*indirect, "load_transactions_per_system_cycle", path); + profile.indirectStoreTransactionsPerCycle = + reader.number(*indirect, "store_transactions_per_system_cycle", path); + profile.indirectDependencyLatencyCycles = + reader.number(*indirect, "dependency_latency_system_cycles", path); + } + if (const auto *control = reader.object(*resources, "control_flow", prefix)) { + const std::string path = prefix + ".control_flow"; + profile.controlFlow.loopBackedgeCycles = + reader.number(*control, "loop_backedge_system_cycles", path); + profile.controlFlow.conditionalBranchCycles = + reader.number(*control, "conditional_branch_system_cycles", path); + profile.controlFlow.divergentBranchPenaltyCycles = + reader.number(*control, "divergent_branch_penalty_system_cycles", path); + profile.controlFlow.synchronizationCycles = + reader.number(*control, "synchronization_system_cycles", path); + } +} + static llvm::Expected loadCandidateProfile(llvm::StringRef requestedPath) { std::string path = requestedPath.empty() ? getDefaultSimdSimtProfilePath() @@ -744,58 +385,20 @@ loadCandidateProfile(llvm::StringRef requestedPath) { const MicrobenchmarkProfile *microbench = microbenchmarkProfile ? &*microbenchmarkProfile : nullptr; - profile.profileVersion = reader.string(*root, "profile_version", "profile"); - const bool usesAnchorPartitionProfile = true; - profile.target = reader.string(*root, "target", "profile"); - if (microbench && llvm::StringRef(profile.target) != microbench->getTarget()) - reader.setError("selection profile target '" + profile.target + + HardwareProfile &hardware = profile.hardware; + hardware.profileVersion = reader.string(*root, "profile_version", "profile"); + hardware.target = reader.string(*root, "target", "profile"); + if (microbench && llvm::StringRef(hardware.target) != microbench->getTarget()) + reader.setError("selection profile target '" + hardware.target + "' does not match shared microbenchmark target '" + microbench->getTarget().str() + "'"); profile.scoreUnit = reader.string(*root, "score_unit", "profile"); - if (const auto *targets = - reader.array(*root, "compatible_targets", "profile")) { - for (const llvm::json::Value &target : *targets) { - if (auto text = target.getAsString()) - profile.compatibleTargets.push_back(text->str()); - else - reader.setError("profile.compatible_targets entries must be strings"); - } - } const auto *calibration = reader.object(*root, "selection_calibration", "profile"); if (calibration) { - profile.programIssueScale = reader.number( - *calibration, "program_issue_scale", "profile.selection_calibration"); - if (const auto *structural = reader.object(*calibration, "simd_structural_penalty_ratio", "profile.selection_calibration")) { - profile.structural.irregularPerDensity = - reader.number(*structural, "irregular_per_density", "structural"); - profile.structural.irregularCap = - reader.number(*structural, "irregular_cap", "structural"); - profile.structural.tinyDotIrregularPerDensity = reader.number( - *structural, "tiny_dot_irregular_per_density", "structural"); - profile.structural.tinyDotIrregularCap = - reader.number(*structural, "tiny_dot_irregular_cap", "structural"); - profile.structural.perMaskRank = - reader.number(*structural, "per_mask_rank", "structural"); - profile.structural.maskCap = - reader.number(*structural, "mask_cap", "structural"); - profile.structural.perWeightedReduction = - reader.number(*structural, "per_weighted_reduction", "structural"); - profile.structural.reductionCap = - reader.number(*structural, "reduction_cap", "structural"); - profile.structural.perStaticLoopTrip = - reader.number(*structural, "per_static_loop_trip", "structural"); - profile.structural.loopCap = - reader.number(*structural, "loop_cap", "structural"); - profile.structural.controlFlow = - reader.number(*structural, "control_flow", "structural"); - profile.structural.rank1IndirectVectorReduction = reader.number( - *structural, "rank1_indirect_vector_reduction", "structural"); - profile.structural.tinyDot = - reader.number(*structural, "tiny_dot", "structural"); profile.structural.tinyDotFlopsMax = reader.integer(*structural, "tiny_dot_flops_max", "structural"); } @@ -804,286 +407,149 @@ loadCandidateProfile(llvm::StringRef requestedPath) { const auto *simd = reader.object(*root, "simd", "profile"); if (simd) { if (simd->getString("vector_width_measurement")) { - profile.simdVectorWidthBits = - static_cast(std::llround(resolveNumberOrMeasurement( - *simd, "vector_width_bits", "vector_width_measurement", "bit", - microbench, reader, "simd"))); + hardware.simd.vectorWidth = std::max( + 1, static_cast(std::llround(resolveNumberOrMeasurement( + *simd, "vector_width_bits", "vector_width_measurement", "bit", + microbench, reader, "simd"))) / + 32); } else { - profile.simdVectorWidthBits = - reader.integer(*simd, "vector_width_bits", "simd"); + hardware.simd.vectorWidth = std::max( + 1, reader.integer(*simd, "vector_width_bits", "simd") / 32); } + hardware.simd.issueWidth = hardware.simd.vectorWidth; if (const auto *startup = reader.object(*simd, "startup_system_cycles", "simd")) - profile.simdSetupCycles = + hardware.simd.setupCycles = reader.number(*startup, "vector", "simd.startup_system_cycles"); if (const auto *ops = reader.object(*simd, "ops", "simd")) { for (llvm::StringRef op : {"f32.add", "f32.sub", "f32.mul", "f32.div", "f32.max", "f32.abs", "f32.exp", "f32.log", "predicate.cmp", "predicate.select", "convert.cast", "f32.clamp"}) - profile.simdOps[op] = resolveOpProfile( + hardware.simd.operationRates[op] = resolveOpProfile( *ops, op, "throughput_vector_instructions_per_system_cycle", "vector_instruction/system_cycle", microbench, reader); } if (const auto *memory = reader.object(*simd, "memory", "simd")) { - profile.simdMte2BytesPerCycle = reader.number( + hardware.simd.loadBytesPerCycle = reader.number( *memory, "vector_mte2_bytes_per_system_cycle", "simd.memory"); - profile.simdMte3BytesPerCycle = + hardware.simd.storeBytesPerCycle = reader.number(*memory, "mte3_bytes_per_system_cycle", "simd.memory"); - profile.simdMemoryConfidence = - reader.optionalString(*memory, "confidence", "none"); } if (const auto *dot = reader.object(*simd, "dot", "simd")) { - profile.simdDotSetupCycles = + hardware.simd.dotSetupCycles = reader.number(*dot, "startup_system_cycles", "simd.dot"); - profile.simdDotFlopsPerCycle = + hardware.simd.dotFlopsPerCycle = reader.number(*dot, "flops_per_system_cycle", "simd.dot"); - profile.simdDotConfidence = - reader.optionalString(*dot, "confidence", "none"); - } - if (const auto *resources = - reader.object(*simd, "stage_resources", "simd")) { - profile.simdStageResources.scalarOperationsPerCycle = - reader.number(*resources, "scalar_operations_per_system_cycle", - "simd.stage_resources"); - profile.simdStageResources.issueOperationsPerCycle = - reader.number(*resources, "issue_instructions_per_system_cycle", - "simd.stage_resources"); - profile.simdStageResources.spillTransactionsPerCycle = - reader.number(*resources, "spill_transactions_per_system_cycle", - "simd.stage_resources"); - if (const auto *indirect = reader.object(*resources, "indirect_memory", - "simd.stage_resources")) { - profile.simdStageResources.indirectLoadTransactionsPerCycle = - reader.number(*indirect, "load_transactions_per_system_cycle", - "simd.stage_resources.indirect_memory"); - profile.simdStageResources.indirectStoreTransactionsPerCycle = - reader.number(*indirect, "store_transactions_per_system_cycle", - "simd.stage_resources.indirect_memory"); - profile.simdStageResources.indirectDependencyLatencyCycles = - reader.number(*indirect, "dependency_latency_system_cycles", - "simd.stage_resources.indirect_memory"); - } - if (const auto *control = reader.object(*resources, "control_flow", - "simd.stage_resources")) { - profile.simdStageResources.controlFlow.loopBackedgeCycles = - reader.number(*control, "loop_backedge_system_cycles", - "simd.stage_resources.control_flow"); - profile.simdStageResources.controlFlow.conditionalBranchCycles = - reader.number(*control, "conditional_branch_system_cycles", - "simd.stage_resources.control_flow"); - profile.simdStageResources.controlFlow.divergentBranchPenaltyCycles = - reader.number(*control, "divergent_branch_penalty_system_cycles", - "simd.stage_resources.control_flow"); - profile.simdStageResources.controlFlow.synchronizationCycles = - reader.number(*control, "synchronization_system_cycles", - "simd.stage_resources.control_flow"); - } } + readStageResources(reader, *simd, "simd", hardware.simd); + const auto predicate = hardware.simd.operationRates.lookup("predicate.cmp"); + hardware.simd.predicateOperationsPerCycle = + predicate.throughput / std::max(1.0, predicate.factor); + hardware.simd.shuffleLanesPerCycle = hardware.simd.vectorWidth; } const auto *simt = reader.object(*root, "simt", "profile"); if (simt) { if (simt->getString("warp_size_measurement")) { - profile.simtWarpSize = + hardware.simt.issueWidth = static_cast(std::llround(resolveNumberOrMeasurement( *simt, "warp_size", "warp_size_measurement", "lane", microbench, reader, "simt"))); } else { - profile.simtWarpSize = reader.integer(*simt, "warp_size", "simt"); + hardware.simt.issueWidth = reader.integer(*simt, "warp_size", "simt"); } + hardware.simt.vectorWidth = 1; if (const auto *setup = reader.object(*simt, "setup_system_cycles", "simt")) { - profile.simtSetupCycles = resolveNumberOrMeasurement( + hardware.simt.setupCycles = resolveNumberOrMeasurement( *setup, "empty_launch", "empty_launch_measurement", "system_cycle", - microbench, reader, "simt.setup_system_cycles", - &profile.simtSetupConfidence); + microbench, reader, "simt.setup_system_cycles"); } if (const auto *ops = reader.object(*simt, "ops", "simt")) { for (llvm::StringRef op : {"f32.add", "f32.sub", "f32.mul", "f32.div", "f32.max", "f32.abs", "f32.exp", "f32.log", "predicate.cmp", "predicate.select", "convert.cast", "f32.clamp"}) - profile.simtOps[op] = + hardware.simt.operationRates[op] = resolveOpProfile(*ops, op, "throughput_scalar_ops_per_system_cycle", "scalar_op/system_cycle", microbench, reader); } if (const auto *dot = reader.object(*simt, "dot", "simt")) { - profile.simtDotSetupCycles = + hardware.simt.dotSetupCycles = reader.number(*dot, "startup_system_cycles", "simt.dot"); - profile.simtDotFlopsPerCycle = + hardware.simt.dotFlopsPerCycle = reader.number(*dot, "flops_per_system_cycle", "simt.dot"); - profile.simtDotConfidence = - reader.optionalString(*dot, "confidence", "none"); - } - if (const auto *camodel = - reader.object(*simt, "camodel_effective", "simt")) { - if (const auto *rates = - reader.object(*camodel, "warp_instructions_per_system_cycle", - "simt.camodel_effective")) - profile.simtPredicateRate = - reader.number(*rates, "predicate", "simt.camodel_effective.rates"); } + const auto predicate = hardware.simt.operationRates.lookup("predicate.cmp"); + hardware.simt.predicateOperationsPerCycle = + predicate.throughput / std::max(1.0, predicate.factor); if (const auto *shuffle = reader.object(*simt, "shuffle", "simt")) { - std::string measuredConfidence; - profile.simtShuffleRate = resolveNumberOrMeasurement( - *shuffle, "warp_instructions_per_system_cycle", - "throughput_measurement", "warp_instruction/system_cycle", microbench, - reader, "simt.shuffle", &measuredConfidence); - profile.simtShuffleConfidence = - reader.optionalString(*shuffle, "confidence", measuredConfidence); + hardware.simt.shuffleLanesPerCycle = + hardware.simt.issueWidth * + resolveNumberOrMeasurement( + *shuffle, "warp_instructions_per_system_cycle", + "throughput_measurement", "warp_instruction/system_cycle", + microbench, reader, "simt.shuffle"); } if (const auto *memory = reader.object(*simt, "memory", "simt")) { - std::string loadConfidence; - std::string storeConfidence; - profile.simtLoadWarpRate = resolveNumberOrMeasurement( + hardware.simt.loadWarpInstructionsPerCycle = resolveNumberOrMeasurement( *memory, "load_warp_instructions_per_system_cycle", "load_throughput_measurement", "warp_instruction/system_cycle", - microbench, reader, "simt.memory", &loadConfidence); - profile.simtStoreWarpRate = resolveNumberOrMeasurement( + microbench, reader, "simt.memory"); + hardware.simt.storeWarpInstructionsPerCycle = resolveNumberOrMeasurement( *memory, "store_warp_instructions_per_system_cycle", "store_throughput_measurement", "warp_instruction/system_cycle", - microbench, reader, "simt.memory", &storeConfidence); - profile.simtMemoryConfidence = reader.optionalString( - *memory, "confidence", - minimumConfidence({loadConfidence, storeConfidence})); - } - if (const auto *resources = - reader.object(*simt, "stage_resources", "simt")) { - profile.simtStageResources.scalarOperationsPerCycle = - reader.number(*resources, "scalar_operations_per_system_cycle", - "simt.stage_resources"); - profile.simtStageResources.issueOperationsPerCycle = - reader.number(*resources, "issue_instructions_per_system_cycle", - "simt.stage_resources"); - profile.simtStageResources.spillTransactionsPerCycle = - reader.number(*resources, "spill_transactions_per_system_cycle", - "simt.stage_resources"); - if (const auto *indirect = reader.object(*resources, "indirect_memory", - "simt.stage_resources")) { - profile.simtStageResources.indirectLoadTransactionsPerCycle = - reader.number(*indirect, "load_transactions_per_system_cycle", - "simt.stage_resources.indirect_memory"); - profile.simtStageResources.indirectStoreTransactionsPerCycle = - reader.number(*indirect, "store_transactions_per_system_cycle", - "simt.stage_resources.indirect_memory"); - profile.simtStageResources.indirectDependencyLatencyCycles = - reader.number(*indirect, "dependency_latency_system_cycles", - "simt.stage_resources.indirect_memory"); - } - if (const auto *control = reader.object(*resources, "control_flow", - "simt.stage_resources")) { - profile.simtStageResources.controlFlow.loopBackedgeCycles = - reader.number(*control, "loop_backedge_system_cycles", - "simt.stage_resources.control_flow"); - profile.simtStageResources.controlFlow.conditionalBranchCycles = - reader.number(*control, "conditional_branch_system_cycles", - "simt.stage_resources.control_flow"); - profile.simtStageResources.controlFlow.divergentBranchPenaltyCycles = - reader.number(*control, "divergent_branch_penalty_system_cycles", - "simt.stage_resources.control_flow"); - profile.simtStageResources.controlFlow.synchronizationCycles = - reader.number(*control, "synchronization_system_cycles", - "simt.stage_resources.control_flow"); - } - if (const auto *superblock = - reader.object(*resources, "superblock", "simt.stage_resources")) { - profile.superblockUsefulFactorLimit = - reader.integer(*superblock, "useful_factor_limit", - "simt.stage_resources.superblock"); - profile.superblockPersistentStatePressureFreeFactor = - reader.integer(*superblock, "persistent_state_pressure_free_factor", - "simt.stage_resources.superblock"); - profile.superblockPersistentStateBytesPerCycle = reader.number( + microbench, reader, "simt.memory"); + } + readStageResources(reader, *simt, "simt", hardware.simt); + if (const auto *resources = simt->getObject("stage_resources")) { + if (const auto *superblock = resources->getObject("superblock")) { + hardware.superblockUsefulFactorLimit = + reader.integer(*superblock, "useful_factor_limit", "superblock"); + hardware.superblockPersistentStatePressureFreeFactor = reader.integer( + *superblock, "persistent_state_pressure_free_factor", "superblock"); + hardware.superblockPersistentStateBytesPerCycle = reader.number( *superblock, "persistent_state_bytes_per_system_cycle", - "simt.stage_resources.superblock"); + "superblock"); } - if (const auto *handoff = reader.object(*resources, "scope_handoff", - "simt.stage_resources")) { - profile.scopeHandoffFixedDirectionalCycles = - reader.number(*handoff, "fixed_directional_system_cycles", - "simt.stage_resources.scope_handoff"); - profile.scopeSimdUbLoadBytesPerCycle = - reader.number(*handoff, "simd_ub_load_bytes_per_system_cycle", - "simt.stage_resources.scope_handoff"); - profile.scopeSimdUbStoreBytesPerCycle = - reader.number(*handoff, "simd_ub_store_bytes_per_system_cycle", - "simt.stage_resources.scope_handoff"); - profile.scopeSimtUbLoadBytesPerThreadPerCycle = reader.number( + if (const auto *handoff = resources->getObject("scope_handoff")) { + hardware.transition.simdToSimtCycles = + hardware.transition.simtToSimdCycles = reader.number( + *handoff, "fixed_directional_system_cycles", "scope_handoff"); + hardware.transition.simdUbLoadBytesPerCycle = reader.number( + *handoff, "simd_ub_load_bytes_per_system_cycle", "scope_handoff"); + hardware.transition.simdUbStoreBytesPerCycle = reader.number( + *handoff, "simd_ub_store_bytes_per_system_cycle", "scope_handoff"); + hardware.transition.simtUbLoadBytesPerThreadPerCycle = reader.number( *handoff, "simt_ub_load_bytes_per_thread_per_system_cycle", - "simt.stage_resources.scope_handoff"); - profile.scopeSimtUbStoreBytesPerThreadPerCycle = reader.number( + "scope_handoff"); + hardware.transition.simtUbStoreBytesPerThreadPerCycle = reader.number( *handoff, "simt_ub_store_bytes_per_thread_per_system_cycle", - "simt.stage_resources.scope_handoff"); + "scope_handoff"); } } - const llvm::json::Object *mixedSetupFallback = nullptr; - if (usesAnchorPartitionProfile) - mixedSetupFallback = reader.object(*simt, "mixed_setup_fallback", "simt"); - else - mixedSetupFallback = reader.object(*simt, "transition", "simt"); - if (mixedSetupFallback) { - for (int64_t numWarps : {1, 2, 4, 8, 16, 32}) { - std::string key = std::to_string(numWarps); - const auto *entry = mixedSetupFallback->getObject(key); - if (!entry) - continue; - std::string measuredConfidence; - profile.mixedSetupFallbacks.push_back( - {numWarps, - resolveNumberOrMeasurement( - *entry, "empty_simt_setup_system_cycles", "measurement", - "system_cycle", microbench, reader, - "simt.mixed_setup_fallback." + key, &measuredConfidence)}); - if (profile.mixedSetupFallbackConfidence == "none") - profile.mixedSetupFallbackConfidence = measuredConfidence; - } - profile.mixedSetupFallbackConfidence = - reader.optionalString(*mixedSetupFallback, "confidence", - profile.mixedSetupFallbackConfidence); - } + hardware.transition.simtWarpSize = hardware.simt.issueWidth; } if (reader.failed()) return llvm::createStringError( std::errc::invalid_argument, "invalid SIMD/SIMT profile '%s': %s", path.c_str(), reader.getError().str().c_str()); - if (profile.profileVersion != "david-v100-simd-simt-20260820-v17") + if (hardware.profileVersion != "david-v100-simd-simt-20260824-v19") return llvm::createStringError( std::errc::invalid_argument, "unsupported SIMD/SIMT profile version '%s' " - "(expected david-v100-simd-simt-20260820-v17)", - profile.profileVersion.c_str()); - const bool usesSharedMicrobench = true; - if (usesSharedMicrobench && !microbench) + "(expected david-v100-simd-simt-20260824-v19)", + hardware.profileVersion.c_str()); + if (!microbench) return llvm::createStringError(std::errc::invalid_argument, - "SIMD/SIMT v17 profile must reference " + "SIMD/SIMT v19 profile must reference " "microbenchmark_profile"); - if (*selectionSchemaVersion != 10) - return llvm::createStringError(std::errc::invalid_argument, - "SIMD/SIMT v17 requires schema_version 10"); - if (profile.simdVectorWidthBits <= 0 || profile.simtWarpSize <= 0 || - profile.simdMte2BytesPerCycle <= 0.0 || - profile.simdMte3BytesPerCycle <= 0.0 || profile.simtLoadWarpRate <= 0.0 || - profile.simtStoreWarpRate <= 0.0 || profile.simtShuffleRate <= 0.0 || - profile.simtPredicateRate <= 0.0 || - profile.superblockUsefulFactorLimit <= 0 || - profile.superblockPersistentStatePressureFreeFactor <= 0 || - profile.superblockPersistentStatePressureFreeFactor > - profile.superblockUsefulFactorLimit || - profile.superblockPersistentStateBytesPerCycle <= 0.0 || - profile.scopeHandoffFixedDirectionalCycles < 0.0 || - profile.scopeSimdUbLoadBytesPerCycle <= 0.0 || - profile.scopeSimdUbStoreBytesPerCycle <= 0.0 || - profile.scopeSimtUbLoadBytesPerThreadPerCycle <= 0.0 || - profile.scopeSimtUbStoreBytesPerThreadPerCycle <= 0.0 || - !profile.simdStageResources.isValid() || - !profile.simtStageResources.isValid() || - profile.mixedSetupFallbacks.empty()) + if (!hardware.isValid()) return llvm::createStringError( std::errc::invalid_argument, - "SIMD/SIMT profile contains non-positive rates or no mixed setup " - "fallbacks"); + "SIMD/SIMT profile contains invalid hardware rates"); std::string canonicalProfile; llvm::raw_string_ostream canonicalStream(canonicalProfile); emitPythonCanonicalJSON(*parsed, canonicalStream); @@ -1110,206 +576,59 @@ loadCandidateProfile(llvm::StringRef requestedPath) { return profile; } -static int64_t mapValue(const llvm::StringMap &values, - llvm::StringRef key, int64_t fallback = 0) { - auto iterator = values.find(key); - return iterator == values.end() ? fallback : iterator->second; -} - -static std::vector> -getProfileOpElements(const SimdSimtFeatureSummary &features) { - const int64_t maxNumel = std::max(1, features.maxTensorNumel); - auto work = [&](llvm::StringRef elementName, int64_t rawCount) { - auto iterator = features.opElements.find(elementName); - if (iterator != features.opElements.end()) - return std::max(0, iterator->second); - return std::max(0, rawCount) * maxNumel; - }; - return { - {"f32.add", work("add", features.addOps)}, - {"f32.sub", work("sub", features.subOps)}, - {"f32.mul", work("mul", features.mulOps)}, - {"f32.div", work("div", features.divOps)}, - {"f32.max", work("max", features.maxOps)}, - {"f32.abs", work("abs", features.absOps)}, - {"f32.exp", work("exp", features.expOps)}, - {"f32.log", work("log", features.logOps)}, - {"predicate.cmp", work("cmp", features.cmpOps)}, - {"predicate.select", work("select", features.selectOps)}, - {"convert.cast", work("cast", features.castOps)}, - {"f32.clamp", work("clamp", features.clampOps)}, - }; -} - -static std::vector> -getProfileOpElements(const SimtAnchorFeatureSummary &features) { - auto work = [&](llvm::StringRef elementName) { - auto iterator = features.opElements.find(elementName); - return iterator == features.opElements.end() - ? int64_t{0} - : std::max(0, iterator->second); - }; - return { - {"f32.add", work("add")}, {"f32.sub", work("sub")}, - {"f32.mul", work("mul")}, {"f32.div", work("div")}, - {"f32.max", work("max")}, {"f32.abs", work("abs")}, - {"f32.exp", work("exp")}, {"f32.log", work("log")}, - {"predicate.cmp", work("cmp")}, {"predicate.select", work("select")}, - {"convert.cast", work("cast")}, {"f32.clamp", work("clamp")}, - }; -} - -static HardwareProfile -buildStageHardwareProfile(const CandidateProfile &profile, unsigned numWarps) { - HardwareProfile hardware; - hardware.profileVersion = profile.profileVersion; - hardware.target = profile.target; - hardware.logicalWarpGroupCount = std::max(1, numWarps); - hardware.superblockUsefulFactorLimit = profile.superblockUsefulFactorLimit; - hardware.superblockPersistentStatePressureFreeFactor = - profile.superblockPersistentStatePressureFreeFactor; - hardware.superblockPersistentStateBytesPerCycle = - profile.superblockPersistentStateBytesPerCycle; - hardware.simd.setupCycles = profile.simdSetupCycles; - hardware.simd.vectorWidth = - std::max(1, profile.simdVectorWidthBits / 32); - hardware.simd.issueWidth = hardware.simd.vectorWidth; - hardware.simt.setupCycles = profile.simtSetupCycles; - hardware.simt.vectorWidth = 1; - hardware.simt.issueWidth = std::max(1, profile.simtWarpSize); - for (const auto &entry : profile.simdOps) - hardware.simd.operationRates[entry.first()] = {entry.second.throughput, - entry.second.factor}; - for (const auto &entry : profile.simtOps) - hardware.simt.operationRates[entry.first()] = {entry.second.throughput, - entry.second.factor}; - hardware.simd.loadBytesPerCycle = profile.simdMte2BytesPerCycle; - hardware.simd.storeBytesPerCycle = profile.simdMte3BytesPerCycle; - hardware.simt.loadWarpInstructionsPerCycle = profile.simtLoadWarpRate; - hardware.simt.storeWarpInstructionsPerCycle = profile.simtStoreWarpRate; - const OpProfile simdPredicate = profile.simdOps.lookup("predicate.cmp"); - const OpProfile simtPredicate = profile.simtOps.lookup("predicate.cmp"); - hardware.simd.predicateOperationsPerCycle = - simdPredicate.throughput / std::max(1.0, simdPredicate.factor); - hardware.simt.predicateOperationsPerCycle = - simtPredicate.throughput / std::max(1.0, simtPredicate.factor); - hardware.simd.shuffleLanesPerCycle = hardware.simd.vectorWidth; - hardware.simt.shuffleLanesPerCycle = - profile.simtWarpSize * profile.simtShuffleRate; - hardware.simd.dotSetupCycles = profile.simdDotSetupCycles; - hardware.simd.dotFlopsPerCycle = profile.simdDotFlopsPerCycle; - hardware.simt.dotSetupCycles = profile.simtDotSetupCycles; - hardware.simt.dotFlopsPerCycle = profile.simtDotFlopsPerCycle; - hardware.simd.scalarOperationsPerCycle = - profile.simdStageResources.scalarOperationsPerCycle; - hardware.simt.scalarOperationsPerCycle = - profile.simtStageResources.scalarOperationsPerCycle; - hardware.simd.issueOperationsPerCycle = - profile.simdStageResources.issueOperationsPerCycle; - hardware.simt.issueOperationsPerCycle = - profile.simtStageResources.issueOperationsPerCycle; - hardware.simd.spillTransactionsPerCycle = - profile.simdStageResources.spillTransactionsPerCycle; - hardware.simt.spillTransactionsPerCycle = - profile.simtStageResources.spillTransactionsPerCycle; - hardware.simd.indirectLoadTransactionsPerCycle = - profile.simdStageResources.indirectLoadTransactionsPerCycle; - hardware.simd.indirectStoreTransactionsPerCycle = - profile.simdStageResources.indirectStoreTransactionsPerCycle; - hardware.simd.indirectDependencyLatencyCycles = - profile.simdStageResources.indirectDependencyLatencyCycles; - hardware.simt.indirectLoadTransactionsPerCycle = - profile.simtStageResources.indirectLoadTransactionsPerCycle; - hardware.simt.indirectStoreTransactionsPerCycle = - profile.simtStageResources.indirectStoreTransactionsPerCycle; - hardware.simt.indirectDependencyLatencyCycles = - profile.simtStageResources.indirectDependencyLatencyCycles; - hardware.simd.controlFlow = profile.simdStageResources.controlFlow; - hardware.simt.controlFlow = profile.simtStageResources.controlFlow; - hardware.transition.simdToSimtCycles = - profile.scopeHandoffFixedDirectionalCycles; - hardware.transition.simtToSimdCycles = - profile.scopeHandoffFixedDirectionalCycles; - hardware.transition.simdUbLoadBytesPerCycle = - profile.scopeSimdUbLoadBytesPerCycle; - hardware.transition.simdUbStoreBytesPerCycle = - profile.scopeSimdUbStoreBytesPerCycle; - hardware.transition.simtUbLoadBytesPerThreadPerCycle = - profile.scopeSimtUbLoadBytesPerThreadPerCycle; - hardware.transition.simtUbStoreBytesPerThreadPerCycle = - profile.scopeSimtUbStoreBytesPerThreadPerCycle; - hardware.transition.simtWarpSize = profile.simtWarpSize; - hardware.transition.source = - "exact scope tensor bytes crossing SIMD/SIMT register files through UB; " - "directional setup remains separate from standalone SIMT VF startup"; - return hardware; -} - static llvm::Expected> evaluateStageModel( const SimdSimtFeatureSummary &features, const CandidateProfile &profile, unsigned numWarps, bool wholeKernelSuperblockMaterializable, - bool scopeSuperblockMaterializable, ModuleOp module = nullptr, - const SimtAnchorPlan *anchorPlan = nullptr) { + bool scopeSuperblockMaterializable, int64_t logicalProgramCountHint, + int64_t physicalCoreCountHint, ModuleOp module, + const SimtAnchorPlan *anchorPlan) { StagePartitionerOptions partitionerOptions; partitionerOptions.tinyDotFlopsMax = profile.structural.tinyDotFlopsMax; partitionerOptions.maximumSuperblockFactor = - (wholeKernelSuperblockMaterializable || features.autoBlockifyV1Applied) + (wholeKernelSuperblockMaterializable || scopeSuperblockMaterializable || + features.autoBlockifyV1Applied) ? 4 : 1; + const int64_t warpLimitedMaximum = numWarps <= 16 ? 4 + : numWarps <= 32 ? 2 + : 1; + partitionerOptions.maximumSuperblockFactor = + std::min(partitionerOptions.maximumSuperblockFactor, warpLimitedMaximum); + if (logicalProgramCountHint > 0) { + const int64_t runtimeMaximum = logicalProgramCountHint >= 4 ? 4 + : logicalProgramCountHint >= 2 ? 2 + : 1; + partitionerOptions.maximumSuperblockFactor = + std::min(partitionerOptions.maximumSuperblockFactor, runtimeMaximum); + } partitionerOptions.scopeSuperblockMaterializable = scopeSuperblockMaterializable; StagePartitioner partitioner; + if (!module || !anchorPlan) + return llvm::createStringError(std::errc::invalid_argument, + "Stage model requires PreparedTTIR and " + "its anchor plan"); auto partition = - module && anchorPlan ? partitioner.partition(module, *anchorPlan, - features, partitionerOptions) - : anchorPlan - ? partitioner.partition(features, partitionerOptions, *anchorPlan) - : partitioner.partition(features, partitionerOptions); + partitioner.partition(module, *anchorPlan, features, partitionerOptions); if (!partition) return partition.takeError(); if (!*partition) return std::optional{}; - HardwareProfile hardwareProfile = - buildStageHardwareProfile(profile, numWarps); - ProfileProvider provider(std::move(hardwareProfile)); - auto snapshot = provider.getSnapshot(profile.target, profile.profileVersion); - if (!snapshot) - return snapshot.takeError(); + HardwareProfile hardwareProfile = profile.hardware; + hardwareProfile.logicalWarpGroupCount = std::max(1, numWarps); StageCostEvaluator evaluator; - auto costTable = evaluator.evaluate(**partition, **snapshot); + auto costTable = evaluator.evaluate(**partition, hardwareProfile); if (!costTable) return costTable.takeError(); - auto routes = solveStageRoutes(*costTable, (*snapshot)->transition); + costTable->logicalProgramCountHint = logicalProgramCountHint; + costTable->physicalCoreCountHint = physicalCoreCountHint; + auto routes = solveStageRoutes(*costTable, hardwareProfile.transition); if (!routes) return routes.takeError(); return std::optional{std::move(*routes)}; } -static SimtApplicabilityResult -evaluateSimtApplicability(const SimdSimtFeatureSummary &features, - bool targetSupported) { - SimtApplicabilityResult result; - result.targetSupported = targetSupported; - result.recognizedAnchorCount = features.simtAnchors.recognizedCount; - result.materializableAnchorCount = features.simtAnchors.count; - result.mechanisms = features.simtAnchors.mechanismKinds; - for (const std::string &kind : features.observedMixedKinds) - appendUnique(result.mechanisms, kind); - llvm::sort(result.mechanisms); - result.mechanismDetected = - result.recognizedAnchorCount > 0 || !result.mechanisms.empty(); - result.materializable = - targetSupported && result.materializableAnchorCount > 0; - if (!result.mechanismDetected) - result.reasons.push_back("no_recognized_simt_mechanism"); - else if (!targetSupported) - result.reasons.push_back("target_does_not_support_simt_materialization"); - else if (result.materializableAnchorCount == 0) - result.reasons.push_back("no_materializable_simt_anchor"); - return result; -} - static llvm::SmallVector> legalCandidates(const SimdSimtCandidateScores &scores, bool allSimdLegal, bool allSimtLegal, bool mixedLegal) { @@ -1336,19 +655,6 @@ static SimdSimtCandidateKind chooseBest(const SimdSimtCandidateScores &scores, .second; } -static SimdSimtCandidateKind -chooseRunnerUp(const SimdSimtCandidateScores &scores, bool allSimdLegal, - bool allSimtLegal, bool mixedLegal, SimdSimtCandidateKind best) { - auto candidates = - legalCandidates(scores, allSimdLegal, allSimtLegal, mixedLegal); - return candidates.size() > 1 ? candidates[1].second : best; -} - -static void sortAndUnique(std::vector &values) { - llvm::sort(values); - values.erase(std::unique(values.begin(), values.end()), values.end()); -} - } // namespace llvm::StringRef @@ -1364,163 +670,37 @@ mlir::ascend::stringifySimdSimtCandidate(SimdSimtCandidateKind candidate) { llvm_unreachable("unknown SIMD/SIMT candidate"); } -double SimdSimtCandidateScores::get(SimdSimtCandidateKind candidate) const { - switch (candidate) { - case SimdSimtCandidateKind::AllSIMD: - return allSimd; - case SimdSimtCandidateKind::AllSIMTOnly: - return allSimtOnly; - case SimdSimtCandidateKind::MixedSIMDSIMT: - return mixedSimdSimt; - } - llvm_unreachable("unknown SIMD/SIMT candidate"); -} - llvm::json::Object SimdSimtCandidateScores::toJSON() const { - llvm::json::Object result; - result[kAllSimd] = allSimd; - result[kAllSimtOnly] = allSimtOnly; - result[kMixedSimdSimt] = mixedSimdSimt; - return result; -} - -static llvm::json::Array toReasonJSON(const std::vector &reasons) { - llvm::json::Array result; - for (const std::string &reason : reasons) - result.push_back(reason); - return result; + return llvm::json::Object{{kAllSimd, allSimd}, + {kAllSimtOnly, allSimtOnly}, + {kMixedSimdSimt, mixedSimdSimt}}; } static llvm::json::Object toLowerabilityJSON(const CandidateLowerability &lowerability) { - llvm::json::Object result; - auto route = [](CandidateLoweringStatus status, - const std::vector &reasons) { - llvm::json::Object entry; - entry["status"] = stringifyCandidateLoweringStatus(status).str(); - entry["reasons"] = toReasonJSON(reasons); - return entry; - }; - result[kAllSimd] = route(lowerability.allSimd, lowerability.allSimdReasons); - result[kAllSimtOnly] = - route(lowerability.allSimtOnly, lowerability.allSimtOnlyReasons); - result[kMixedSimdSimt] = route(lowerability.mixed, lowerability.mixedReasons); - return result; -} - -static llvm::json::Object toAtomicFactsJSON(const TensorAtomicFacts &facts) { - llvm::json::Object result; - result["update_elements"] = facts.updateElements; - result["address_rank"] = facts.addressRank; - result["value_type"] = facts.valueType; - result["offset_type"] = facts.offsetType; - result["operation"] = facts.operation; - result["has_mask"] = facts.hasMask; - if (facts.staticMaskActiveFraction) - result["static_mask_active_fraction"] = *facts.staticMaskActiveFraction; - else - result["static_mask_active_fraction"] = nullptr; - result["result_used"] = facts.resultUsed; - result["address_is_lane_varying"] = facts.addressIsLaneVarying; - result["address_depends_on_loaded_index"] = facts.addressDependsOnLoadedIndex; - result["contention"] = facts.contention; - return result; -} - -static llvm::json::Object toHistogramFactsJSON(const HistogramFacts &facts) { - llvm::json::Object result; - result["input_elements"] = facts.inputElements; - result["num_bins"] = facts.numBins; - result["input_type"] = facts.inputType; - result["result_type"] = facts.resultType; - return result; -} - -static llvm::json::Object -toPlainCumsumFactsJSON(const PlainCumsumFacts &facts) { - llvm::json::Object result; - result["axis_extent"] = facts.axisExtent; - result["element_type"] = facts.elementType; - result["reverse"] = facts.reverse; - return result; + return llvm::json::Object{{kAllSimd, lowerability.allSimd}, + {kAllSimtOnly, lowerability.allSimtOnly}, + {kMixedSimdSimt, lowerability.mixed}}; } static llvm::json::Object toTriangularSolveFactsJSON(const TriangularSolveFacts &facts) { - llvm::json::Object result; - result["block_rows"] = facts.blockRows; - result["block_columns"] = facts.blockColumns; - result["accumulator_type"] = facts.accumulatorType; - result["recurrence_start_row"] = facts.recurrenceStartRow; - result["recurrence_loop_count"] = facts.recurrenceLoopCount; - result["dense_dot_tail_ops"] = facts.denseDotTailOps; - result["requires_cube_tail_partition"] = facts.requiresCubeTailPartition; - return result; + return llvm::json::Object{ + {"block_rows", facts.blockRows}, + {"block_columns", facts.blockColumns}, + {"accumulator_type", facts.accumulatorType}, + {"recurrence_start_row", facts.recurrenceStartRow}, + {"recurrence_loop_count", facts.recurrenceLoopCount}, + {"dense_dot_tail_ops", facts.denseDotTailOps}, + {"requires_cube_tail_partition", facts.requiresCubeTailPartition}}; } llvm::json::Object SimtAnchorFeatureSummary::toJSON() const { llvm::json::Object result; - result["recognized_count"] = recognizedCount; result["count"] = count; - result["materializable_count"] = count; - result["covered_operation_count"] = coveredOperationCount; - result["load_ops"] = loadOps; - result["store_ops"] = storeOps; - result["reduce_ops"] = reduceOps; - result["scan_ops"] = scanOps; - result["gather_ops"] = gatherOps; - result["dot_ops"] = dotOps; - result["atomic_ops"] = atomicOps; - result["histogram_ops"] = histogramOps; - result["max_tensor_numel"] = maxTensorNumel; - result["max_element_bits"] = maxElementBits; - result["mask_rank_sum"] = maskRankSum; - result["unique_mask_values"] = uniqueMaskValues; - result["unique_mask_rank_sum"] = uniqueMaskRankSum; - result["predicate_elements"] = predicateElements; - result["predicate_lane_evaluations"] = predicateLaneEvaluations; - result["pointer_tensor_ops"] = pointerTensorOps; - result["loaded_index_dependent_memory_ops"] = loadedIndexDependentMemoryOps; - result["lane_dependent_pointer_ops"] = laneDependentPointerOps; - result["max_reduce_axis_extent"] = maxReduceAxisExtent; - result["weighted_reduce_axis_elements"] = weightedReduceAxisElements; - result["shuffle_lane_steps"] = shuffleLaneSteps; - result["static_loop_count"] = staticLoopCount; - result["static_loop_trip_count_sum"] = staticLoopTripCountSum; - result["modeled_dynamic_loop_count"] = modeledDynamicLoopCount; - result["modeled_dynamic_loop_trip_count_sum"] = - modeledDynamicLoopTripCountSum; result["conditional_branch_count"] = conditionalBranchCount; result["divergent_branch_count"] = divergentBranchCount; result["active_lane_ratio"] = activeLaneRatio; - result["has_control_flow"] = hasControlFlow; - result["weighted_ops"] = ::toJSON(weightedOps); - result["op_elements"] = ::toJSON(opElements); - result["load_bytes"] = loadBytes; - result["store_bytes"] = storeBytes; - result["load_warp_instructions"] = loadWarpInstructions; - result["store_warp_instructions"] = storeWarpInstructions; - result["dot_flops"] = dotFlops; - result["captured_tensor_count"] = capturedTensorCount; - result["escaping_tensor_count"] = escapingTensorCount; - result["captured_tensor_bytes"] = capturedTensorBytes; - result["escaping_tensor_bytes"] = escapingTensorBytes; - llvm::json::Array mechanisms; - for (const std::string &kind : mechanismKinds) - mechanisms.push_back(kind); - result["mechanism_kinds"] = std::move(mechanisms); - llvm::json::Array atomicFacts; - for (const TensorAtomicFacts &facts : tensorAtomics) - atomicFacts.push_back(toAtomicFactsJSON(facts)); - result["tensor_atomics"] = std::move(atomicFacts); - llvm::json::Array histogramFacts; - for (const HistogramFacts &facts : histograms) - histogramFacts.push_back(toHistogramFactsJSON(facts)); - result["histograms"] = std::move(histogramFacts); - llvm::json::Array cumsumFacts; - for (const PlainCumsumFacts &facts : plainCumsums) - cumsumFacts.push_back(toPlainCumsumFactsJSON(facts)); - result["plain_cumsums"] = std::move(cumsumFacts); llvm::json::Array triangularFacts; for (const TriangularSolveFacts &facts : triangularSolves) triangularFacts.push_back(toTriangularSolveFactsJSON(facts)); @@ -1534,212 +714,19 @@ llvm::json::Object SimdSimtFeatureSummary::toJSON() const { result["load_ops"] = loadOps; result["store_ops"] = storeOps; result["reduce_ops"] = reduceOps; - result["scan_ops"] = scanOps; - result["gather_ops"] = gatherOps; result["dot_ops"] = dotOps; - result["atomic_ops"] = atomicOps; - result["histogram_ops"] = histogramOps; - result["broadcast_ops"] = broadcastOps; - result["expand_dims_ops"] = expandDimsOps; - result["splat_ops"] = splatOps; - result["addptr_ops"] = addPtrOps; - result["arith_ops"] = arithOps; - result["math_ops"] = mathOps; - result["add_ops"] = addOps; - result["sub_ops"] = subOps; - result["mul_ops"] = mulOps; - result["div_ops"] = divOps; - result["max_ops"] = maxOps; - result["abs_ops"] = absOps; - result["exp_ops"] = expOps; - result["log_ops"] = logOps; - result["cmp_ops"] = cmpOps; - result["select_ops"] = selectOps; - result["cast_ops"] = castOps; - result["clamp_ops"] = clampOps; - result["scalar_ops"] = scalarOps; - result["max_tensor_rank"] = maxTensorRank; - result["max_tensor_numel"] = maxTensorNumel; - result["max_element_bits"] = maxElementBits; - result["mask_tensor_ops"] = maskTensorOps; - result["mask_rank_sum"] = maskRankSum; - result["unique_mask_values"] = uniqueMaskValues; - result["unique_mask_rank_sum"] = uniqueMaskRankSum; - result["predicate_elements"] = predicateElements; - result["predicate_lane_evaluations"] = predicateLaneEvaluations; - result["mask_broadcast_ops"] = maskBroadcastOps; - result["pointer_tensor_ops"] = pointerTensorOps; - result["pointer_unstructured_dims"] = pointerUnstructuredDims; result["loaded_index_dependent_memory_ops"] = loadedIndexDependentMemoryOps; - result["lane_dependent_pointer_ops"] = laneDependentPointerOps; - result["row_local_reduce_ops"] = rowLocalReduceOps; - result["max_reduce_axis_extent"] = maxReduceAxisExtent; - result["weighted_reduce_axis_elements"] = weightedReduceAxisElements; - result["shuffle_lane_steps"] = shuffleLaneSteps; - result["scalar_load_ops"] = scalarLoadOps; - result["scalar_store_ops"] = scalarStoreOps; - result["vector_ptr_splat_ops"] = vectorPtrSplatOps; - result["vector_reduce_to_scalar_ops"] = vectorReduceToScalarOps; - result["rank1_indirect_vector_reduce"] = rank1IndirectVectorReduce; - result["weighted_ops"] = ::toJSON(weightedOps); - result["op_elements"] = ::toJSON(opElements); - result["load_bytes"] = loadBytes; - result["store_bytes"] = storeBytes; - result["load_warp_instructions"] = loadWarpInstructions; - result["store_warp_instructions"] = storeWarpInstructions; result["dot_flops"] = dotFlops; - result["dot_output_elements"] = dotOutputElements; - llvm::json::Array dotShapes; - for (const auto &shape : dotMNK) - dotShapes.push_back(llvm::json::Array({shape[0], shape[1], shape[2]})); - result["dot_mnk"] = std::move(dotShapes); - result["static_loop_count"] = staticLoopCount; - result["static_loop_trip_count_sum"] = staticLoopTripCountSum; result["static_loop_trip_count_max"] = staticLoopTripCountMax; - result["modeled_dynamic_loop_count"] = modeledDynamicLoopCount; - result["modeled_dynamic_loop_trip_count_sum"] = - modeledDynamicLoopTripCountSum; - result["loop_carried_data_dependency_count"] = loopCarriedDataDependencyCount; - result["pointer_induction_dependency_count"] = - pointerInductionDependencyCount; result["conditional_branch_count"] = conditionalBranchCount; result["divergent_branch_count"] = divergentBranchCount; result["active_lane_ratio"] = activeLaneRatio; llvm::json::Object postTransform; - postTransform["ttir_layout_merge_applied"] = ttirLayoutMergeApplied; - postTransform["coalesce_factor"] = coalesceFactor; - postTransform["coalesce_axis"] = coalesceAxis; postTransform["auto_blockify_v1_applied"] = autoBlockifyV1Applied; postTransform["auto_blockify_v1_loop_count"] = autoBlockifyV1LoopCount; - postTransform["auto_blockify_v1_schedule_op_count"] = - autoBlockifyV1ScheduleOpCount; - postTransform["auto_blockify_v1_dynamic_trip_count"] = - autoBlockifyV1HasDynamicTripCount; result["post_transform"] = std::move(postTransform); - result["has_dot"] = hasDot; - result["has_gather"] = hasGather; - result["has_atomic"] = hasAtomic; - result["has_histogram"] = hasHistogram; - result["has_scan"] = hasScan; result["has_explicit_scope"] = hasExplicitScope; - result["has_control_flow"] = hasControlFlow; - result["has_dynamic_shape"] = hasDynamicShape; - result["has_unknown_trip_count"] = hasUnknownTripCount; result["simt_anchors"] = simtAnchors.toJSON(); - llvm::json::Array mixedKinds; - for (const std::string &kind : observedMixedKinds) - mixedKinds.push_back(kind); - result["observed_mixed_kinds"] = std::move(mixedKinds); - result["mixed_required"] = !observedMixedKinds.empty(); - result["mandatory_mixed_enabled"] = false; - return result; -} - -llvm::json::Object SimtApplicabilityResult::toJSON() const { - llvm::json::Object result; - result["mechanism_detected"] = mechanismDetected; - result["target_supported"] = targetSupported; - result["materializable"] = materializable; - result["recognized_anchor_count"] = recognizedAnchorCount; - result["materializable_anchor_count"] = materializableAnchorCount; - llvm::json::Array mechanismValues; - for (const std::string &mechanism : mechanisms) - mechanismValues.push_back(mechanism); - result["mechanisms"] = std::move(mechanismValues); - llvm::json::Array reasonValues; - for (const std::string &reason : reasons) - reasonValues.push_back(reason); - result["reasons"] = std::move(reasonValues); - return result; -} - -llvm::json::Object -SimdSimtCostBreakdown::toJSON(const SimdSimtFeatureSummary &features) const { - llvm::json::Object result; - llvm::json::Object compute; - compute["simd"] = simdComputeCycles; - compute["simt"] = simtComputeCycles; - compute["simd_dot"] = simdDotCycles; - compute["simt_dot"] = simtDotCycles; - result["compute_only"] = std::move(compute); - - llvm::json::Object memory; - memory["load_bytes"] = features.loadBytes; - memory["store_bytes"] = features.storeBytes; - memory["simd_load_system_cycles"] = simdLoadCycles; - memory["simd_store_system_cycles"] = simdStoreCycles; - memory["simd_roofline_system_cycles"] = simdMemoryCycles; - memory["simt_load_warp_instructions"] = features.loadWarpInstructions; - memory["simt_store_warp_instructions"] = features.storeWarpInstructions; - memory["simt_load_system_cycles"] = simtLoadCycles; - memory["simt_store_system_cycles"] = simtStoreCycles; - memory["simt_serial_memory_system_cycles"] = simtMemoryCycles; - // Compatibility field retained for existing report consumers. SIMT memory - // is not roofline-overlapped with SIMT compute by the route model. - memory["simt_roofline_system_cycles"] = simtMemoryCycles; - result["memory"] = std::move(memory); - - llvm::json::Object structure; - structure["irregular_density"] = irregularDensity; - structure["tiny_dot_underfill"] = tinyDotUnderfill; - structure["components"] = ::toJSON(structuralComponents); - structure["penalty_ratio"] = structuralPenaltyRatio; - structure["simd_structural_penalty_system_cycles"] = - simdStructuralPenaltyCycles; - result["structure"] = std::move(structure); - - llvm::json::Object mixed; - mixed["derived_simd_fraction"] = mixedSimdFraction; - mixed["cost_source"] = mixedCostSource; - mixed["setup_fallback_num_warps"] = mixedSetupFallbackNumWarps; - mixed["mixed_setup_fallback_system_cycles"] = mixedSetupFallbackCycles; - mixed["standalone_serialized_setup_system_cycles"] = - standaloneSimtSetupCycles; - mixed["setup_proxy_delta_system_cycles"] = setupProxyDeltaCycles; - mixed["directional_transition_system_cycles"] = nullptr; - mixed["directional_transition_measurement_status"] = "unmeasured"; - llvm::json::Object partition; - partition["simd_regular_compute_system_cycles"] = - mixedSimdRegularComputeCycles; - partition["simd_regular_dot_system_cycles"] = mixedSimdRegularDotCycles; - partition["simd_regular_memory_system_cycles"] = mixedSimdRegularMemoryCycles; - partition["simd_regular_payload_system_cycles"] = - mixedSimdRegularPayloadCycles; - partition["simt_anchor_compute_system_cycles"] = mixedSimtAnchorComputeCycles; - partition["simt_anchor_dot_system_cycles"] = mixedSimtAnchorDotCycles; - partition["simt_anchor_memory_system_cycles"] = mixedSimtAnchorMemoryCycles; - partition["simt_anchor_shuffle_system_cycles"] = mixedSimtAnchorShuffleCycles; - partition["simt_anchor_predicate_system_cycles"] = - mixedSimtAnchorPredicateCycles; - partition["simt_anchor_payload_system_cycles"] = mixedSimtAnchorPayloadCycles; - partition["simt_anchor_calibrated_payload_system_cycles"] = - mixedSimtAnchorCalibratedPayloadCycles; - partition["cube_tail_dot_ops"] = cubeTailDotOps; - partition["cube_tail_dot_flops"] = cubeTailDotFlops; - partition["measured_boundary_system_cycles"] = nullptr; - partition["applied_boundary_fallback_system_cycles"] = mixedBoundaryCycles; - partition["remaining_simd_structural_penalty_ratio"] = - mixedRemainingStructuralPenaltyRatio; - mixed["partition"] = std::move(partition); - result["mixed"] = std::move(mixed); - - llvm::json::Object execution; - execution["shuffle_warp_instructions"] = simtShuffleInstructions; - execution["shuffle_system_cycles"] = simtShuffleCycles; - execution["predicate_warp_instructions"] = simtPredicateInstructions; - execution["predicate_system_cycles"] = simtPredicateCycles; - execution["program_issue_scale"] = programIssueScale; - execution["simd_setup_system_cycles"] = simdSetupCycles; - execution["simt_setup_system_cycles"] = simtSetupCycles; - execution["simd_issue_payload_system_cycles"] = simdIssuePayloadCycles; - execution["simt_issue_payload_system_cycles"] = simtIssuePayloadCycles; - execution["simt_issue_aggregation"] = "serial_sum"; - result["simt_execution"] = std::move(execution); - - llvm::json::Object opBreakdown; - opBreakdown["simd_ops_system_cycles"] = ::toJSON(simdOpSystemCycles); - opBreakdown["simt_ops_system_cycles"] = ::toJSON(simtOpSystemCycles); - result["op_breakdown"] = std::move(opBreakdown); return result; } @@ -1750,7 +737,6 @@ llvm::json::Object SimdSimtCostReport::toJSON() const { result["profile_version"] = profileVersion; result["profile_target"] = profileTarget; result["actual_target"] = actualTarget; - result["target_compatible"] = targetCompatible; result["profile_content_sha256"] = profileContentSha256; result["selection_profile_content_sha256"] = selectionProfileContentSha256; llvm::json::Object sharedEvidence; @@ -1759,12 +745,8 @@ llvm::json::Object SimdSimtCostReport::toJSON() const { sharedEvidence["content_sha256"] = microbenchmarkProfileContentSha256; result["shared_microbenchmark_profile"] = std::move(sharedEvidence); result["unit"] = scoreUnit; - result["score_scope"] = scoreScope; - result["excludes"] = llvm::json::Array({"host_launch", "grid_wave_count"}); result["candidate_costs"] = candidateCosts.toJSON(); - result["candidate_ratios_to_best"] = candidateRatiosToBest.toJSON(); result["decision_kind"] = stringifySimdSimtCandidate(decision); - result["best_score"] = bestScore; llvm::json::Array selectableCandidates; if (allSimdCandidateLegal) selectableCandidates.push_back(kAllSimd); @@ -1777,26 +759,8 @@ llvm::json::Object SimdSimtCostReport::toJSON() const { for (const std::string &value : unsupported) unsupportedValues.push_back(value); result["unmodeled_cost_terms"] = std::move(unsupportedValues); - result["applicability"] = applicability.toJSON(); result["stage_model"] = stageModel.toJSON(); - llvm::json::Object roles; - roles[kAllSimd] = - allSimdCandidateLegal ? "selectable_candidate" : "inapplicable"; - roles[kAllSimtOnly] = - allSimtOnlyCandidateLegal ? "selectable_candidate" : "inapplicable"; - roles[kMixedSimdSimt] = - mixedCandidateLegal ? "selectable_candidate" : "inapplicable"; - result["candidate_roles"] = std::move(roles); - - llvm::json::Object analytical; - analytical[kAllSimd] = breakdown.simdAnalyticalCycles; - analytical[kAllSimtOnly] = breakdown.simtAnalyticalCycles; - result["analytical_candidate_costs"] = std::move(analytical); - - llvm::json::Object detail = breakdown.toJSON(features); - for (auto &entry : detail) - result[entry.first] = std::move(entry.second); if (includeFeaturesInJSON) result["features"] = features.toJSON(); return result; @@ -1838,556 +802,86 @@ mlir::ascend::analyzeSimdSimtFeatures(ModuleOp module, "cannot analyze a null ModuleOp"); SimdSimtFeatureSummary features; - initializeWorkMaps(features); - initializeWorkMaps(features.simtAnchors); - features.ttirLayoutMergeApplied = - module->hasAttr("ta.ttir_layout_merge.applied"); - if (auto factor = module->getAttrOfType("hacc.coalesce_factor")) - features.coalesceFactor = std::max(1, factor.getInt()); - if (auto axis = module->getAttrOfType("hacc.coalesce_axis")) - features.coalesceAxis = axis.getInt(); - module.walk([&](Operation *op) { - if (op->hasAttr("ta.auto_blockify_v1")) - features.autoBlockifyV1Applied = true; - if (op->hasAttr("ta.auto_blockify_v1.schedule")) - ++features.autoBlockifyV1ScheduleOpCount; - if (!op->hasAttr("ta.auto_blockify_v1.loop")) - return; - features.autoBlockifyV1Applied = true; - ++features.autoBlockifyV1LoopCount; - if (!getKnownStaticLoopTripCount(op)) - features.autoBlockifyV1HasDynamicTripCount = true; - }); - llvm::DenseSet anchorSet; + llvm::DenseSet anchorRoots; llvm::DenseMap structuralTripEstimates; + features.simtAnchors.count = llvm::count_if( + anchorPlan.anchors, + [](const SimtAnchorDescriptor &anchor) { return anchor.materializable; }); + features.simtAnchors.kernelLowerability = anchorPlan.kernelLowerability; + for (const SimtAnchorDescriptor &anchor : anchorPlan.anchors) { + if (anchor.triangularSolve) + features.simtAnchors.triangularSolves.push_back(*anchor.triangularSolve); if (!anchor.materializable) continue; - if (anchor.scopeOperations.empty()) { - if (anchor.operation) - anchorSet.insert(anchor.operation); - } else { - for (Operation *scopeOperation : anchor.scopeOperations) - if (scopeOperation) - anchorSet.insert(scopeOperation); - } - if (anchor.kind == SimtAnchorKind::TriangularSolveLoop) { - // A recognized solve_tril block has a fixed 16x16 state and starts its - // recurrence at row 2: a full block performs 16 - 2 = 14 iterations. - // The TTIR upper bound is min(runtime_remaining, block_end), so it is - // not a compile-time constant even though the full-tile estimate is - // structurally known. Do not apply this fallback to generic loops. - for (Operation *scopeOperation : anchor.scopeOperations) - if (scopeOperation && - scopeOperation->getName().getStringRef() == "scf.for") - structuralTripEstimates[scopeOperation] = 14; - } - } - features.simtAnchors.recognizedCount = anchorPlan.anchors.size(); - features.simtAnchors.count = anchorPlan.materializableCount(); - features.simtAnchors.kernelLowerability = anchorPlan.kernelLowerability; - for (const SimtAnchorDescriptor &anchor : anchorPlan.anchors) { - std::string kind = stringifySimtAnchorKind(anchor.kind).str(); - appendUnique(features.simtAnchors.mechanismKinds, kind); - appendUnique(features.observedMixedKinds, kind); - if (const auto *facts = std::get_if(&anchor.facts)) - features.simtAnchors.tensorAtomics.push_back(*facts); - else if (const auto *facts = std::get_if(&anchor.facts)) - features.simtAnchors.histograms.push_back(*facts); - else if (const auto *facts = std::get_if(&anchor.facts)) - features.simtAnchors.plainCumsums.push_back(*facts); - else if (const auto *facts = - std::get_if(&anchor.facts)) - features.simtAnchors.triangularSolves.push_back(*facts); - } - - auto isInAnchor = [&](Operation *op) { - for (Operation *current = op; current; current = current->getParentOp()) - if (anchorSet.contains(current)) + for (Operation *operation : anchor.scopeOperations) + if (operation) + anchorRoots.insert(operation); + if (anchor.scopeOperations.empty() && anchor.operation) + anchorRoots.insert(anchor.operation); + if (anchor.kind == SimtAnchorKind::TriangularSolveLoop) + for (Operation *operation : anchor.scopeOperations) + if (operation && operation->getName().getStringRef() == "scf.for") + structuralTripEstimates[operation] = 14; + } + + auto isInAnchor = [&](Operation *operation) { + for (; operation; operation = operation->getParentOp()) + if (anchorRoots.contains(operation)) return true; return false; }; - - llvm::DenseSet capturedTensors; - llvm::DenseSet escapingTensors; - llvm::DenseSet uniqueMasks; - llvm::DenseSet anchorUniqueMasks; - auto isValueDefinedInAnchor = [&](Value value) { - if (Operation *definingOp = value.getDefiningOp()) - return isInAnchor(definingOp); - auto argument = dyn_cast(value); - Operation *parent = argument ? argument.getOwner()->getParentOp() : nullptr; - return parent && isInAnchor(parent); - }; - module.walk([&](Operation *op) { - if (!isInAnchor(op)) - return; - for (Value operand : op->getOperands()) { - if (!isa(operand.getType()) || - isValueDefinedInAnchor(operand) || - !capturedTensors.insert(operand).second) - continue; - ++features.simtAnchors.capturedTensorCount; - features.simtAnchors.capturedTensorBytes += - getStaticTensorBytes(operand.getType()); - } - for (Value result : op->getResults()) { - if (!isa(result.getType()) || result.use_empty()) - continue; - bool escapes = llvm::any_of(result.getUses(), [&](OpOperand &use) { - return !isInAnchor(use.getOwner()); - }); - if (!escapes || !escapingTensors.insert(result).second) - continue; - ++features.simtAnchors.escapingTensorCount; - features.simtAnchors.escapingTensorBytes += - getStaticTensorBytes(result.getType()); - } - }); - auto updateTypeStats = [&](Type type, bool inAnchor) { - if (auto tensor = dyn_cast(type)) { - if (!tensor.hasStaticShape()) - features.hasDynamicShape = true; - features.maxElementBits = - std::max(features.maxElementBits, getTypeBitWidth(type)); - features.maxTensorRank = - std::max(features.maxTensorRank, tensor.getRank()); - features.maxTensorNumel = - std::max(features.maxTensorNumel, getStaticNumElements(type)); - if (inAnchor) { - features.simtAnchors.maxElementBits = std::max( - features.simtAnchors.maxElementBits, getTypeBitWidth(type)); - features.simtAnchors.maxTensorNumel = std::max( - features.simtAnchors.maxTensorNumel, getStaticNumElements(type)); - } + module.walk([&](Operation *operation) { + if (operation->hasAttr("ta.auto_blockify_v1")) + features.autoBlockifyV1Applied = true; + if (operation->hasAttr("ta.auto_blockify_v1.loop")) { + features.autoBlockifyV1Applied = true; + ++features.autoBlockifyV1LoopCount; } - }; - - module.walk([&](Operation *op) { - // V1 scheduling is a separate dispatch phase. Keep it in post_transform - // diagnostics, but do not reinterpret its dynamic physical-core loop as - // an unknown-trip algorithm loop or charge its scalar prologue at - // candidate-specific SIMD/SIMT rates. - if (op->hasAttr("ta.auto_blockify_v1.schedule")) + if (operation->hasAttr("ta.auto_blockify_v1.schedule")) return; - llvm::StringRef name = op->getName().getStringRef(); - const int64_t elements = getOperationElements(op); - const int64_t loopMultiplier = - getLoopMultiplier(op, structuralTripEstimates); - const bool inAnchor = isInAnchor(op); - if (inAnchor) - ++features.simtAnchors.coveredOperationCount; - for (Type type : op->getOperandTypes()) - updateTypeStats(type, inAnchor); - for (Type type : op->getResultTypes()) - updateTypeStats(type, inAnchor); - for (Region ®ion : op->getRegions()) - for (Block &block : region) - for (BlockArgument argument : block.getArguments()) - updateTypeStats(argument.getType(), inAnchor); + llvm::StringRef name = operation->getName().getStringRef(); + const bool inAnchor = isInAnchor(operation); + const int64_t multiplier = + getLoopMultiplier(operation, structuralTripEstimates); + features.hasExplicitScope |= name == "scope.scope"; + features.loadOps += name == "tt.load"; + features.storeOps += name == "tt.store"; + features.reduceOps += name == "tt.reduce"; + features.dotOps += name == "tt.dot"; - if (name.starts_with("arith.")) - ++features.arithOps; - if (name.starts_with("math.")) - ++features.mathOps; - if (name.starts_with("scf.") || name.starts_with("cf.")) { - features.hasControlFlow = true; - if (inAnchor) - features.simtAnchors.hasControlFlow = true; - } if (name == "scf.if" || name == "cf.cond_br") { ++features.conditionalBranchCount; if (inAnchor) ++features.simtAnchors.conditionalBranchCount; - const bool laneVarying = - op->getNumOperands() > 0 && - isa(op->getOperand(0).getType()); - if (laneVarying) { - ++features.divergentBranchCount; - if (inAnchor) - ++features.simtAnchors.divergentBranchCount; - } - } - if (name == "scope.scope") - features.hasExplicitScope = true; - - auto incrementRaw = [&](int64_t &counter, int64_t &anchorCounter) { - ++counter; + const bool divergent = + operation->getNumOperands() > 0 && + isa(operation->getOperand(0).getType()); + features.divergentBranchCount += divergent; if (inAnchor) - ++anchorCounter; - }; - if (name == "tt.load") - incrementRaw(features.loadOps, features.simtAnchors.loadOps); - else if (name == "tt.store") - incrementRaw(features.storeOps, features.simtAnchors.storeOps); - else if (name == "tt.reduce") - incrementRaw(features.reduceOps, features.simtAnchors.reduceOps); - else if (name == "tt.scan" || name == "tt.associative_scan") - incrementRaw(features.scanOps, features.simtAnchors.scanOps); - else if (name == "tt.gather") - incrementRaw(features.gatherOps, features.simtAnchors.gatherOps); - else if (name == "tt.dot") - incrementRaw(features.dotOps, features.simtAnchors.dotOps); - else if (name.starts_with("tt.atomic")) - incrementRaw(features.atomicOps, features.simtAnchors.atomicOps); - else if (name == "tt.histogram") - incrementRaw(features.histogramOps, features.simtAnchors.histogramOps); - else if (name == "tt.broadcast") - ++features.broadcastOps; - else if (name == "tt.expand_dims") - ++features.expandDimsOps; - else if (name == "tt.splat") - ++features.splatOps; - else if (name == "tt.addptr") - ++features.addPtrOps; - - if (name == "arith.addf" || name == "arith.addi") - ++features.addOps; - else if (name == "arith.subf" || name == "arith.subi") - ++features.subOps; - else if (name == "arith.mulf" || name == "arith.muli") - ++features.mulOps; - else if (name == "arith.divf" || name == "arith.divsi" || - name == "arith.divui") - ++features.divOps; - else if (name == "arith.maxnumf" || name == "arith.maxf" || - name == "arith.maxsi" || name == "arith.maxui") - ++features.maxOps; - else if (name == "math.absf" || name == "math.absi") - ++features.absOps; - else if (name == "math.exp") - ++features.expOps; - else if (name == "math.log") - ++features.logOps; - else if (name == "arith.cmpf" || name == "arith.cmpi") - ++features.cmpOps; - else if (name == "arith.select") - ++features.selectOps; - else if (isCastOp(name)) - ++features.castOps; - else if (name.starts_with("tt.clamp")) - ++features.clampOps; - - llvm::StringRef weightedKind = classifyWeightedOp(name); - if (!weightedKind.empty()) { - int64_t weightedElements = elements; - if (name == "tt.histogram" && op->getNumOperands() > 0) - if (auto input = - dyn_cast(op->getOperand(0).getType())) - if (input.hasStaticShape()) - weightedElements = getStaticNumElements(input); - features.weightedOps[weightedKind] += loopMultiplier; - features.opElements[weightedKind] += weightedElements * loopMultiplier; - if (inAnchor) { - features.simtAnchors.weightedOps[weightedKind] += loopMultiplier; - features.simtAnchors.opElements[weightedKind] += - weightedElements * loopMultiplier; - } + features.simtAnchors.divergentBranchCount += divergent; } - if (name == "scf.for") { - auto knownTripCount = getKnownStaticLoopTripCount(op); - if (!knownTripCount) - features.hasUnknownTripCount = true; - int64_t tripCount = getModeledLoopTripCount(op, structuralTripEstimates); - const bool usedStructuralEstimate = - !knownTripCount && structuralTripEstimates.contains(op); - ++features.staticLoopCount; - features.staticLoopTripCountSum += tripCount; + const int64_t trip = + getModeledLoopTripCount(operation, structuralTripEstimates); features.staticLoopTripCountMax = - std::max(features.staticLoopTripCountMax, tripCount); - if (usedStructuralEstimate) { - ++features.modeledDynamicLoopCount; - features.modeledDynamicLoopTripCountSum += tripCount; - } - if (inAnchor) { - ++features.simtAnchors.staticLoopCount; - features.simtAnchors.staticLoopTripCountSum += tripCount; - if (usedStructuralEstimate) { - ++features.simtAnchors.modeledDynamicLoopCount; - features.simtAnchors.modeledDynamicLoopTripCountSum += tripCount; - } - } - if (op->getNumRegions() > 0 && !op->getRegion(0).empty()) { - Block &body = op->getRegion(0).front(); - // scf.for block argument 0 is the induction variable; remaining - // arguments are loop-carried iter_args. - for (unsigned argumentIndex = 1; argumentIndex < body.getNumArguments(); - ++argumentIndex) { - BlockArgument argument = body.getArgument(argumentIndex); - if (argument.use_empty()) - continue; - if (isPointerType(argument.getType()) || - isAddressOnlyLoopCarriedValue(argument)) - ++features.pointerInductionDependencyCount; - else - ++features.loopCarriedDataDependencyCount; - } - } + std::max(features.staticLoopTripCountMax, trip); } - - auto dataTypeAndElements = [&](bool load) -> std::pair { - if (load && op->getNumResults() > 0) - return {op->getResult(0).getType(), - getStaticNumElements(op->getResult(0).getType())}; - if (!load && op->getNumOperands() > 1) - return {op->getOperand(1).getType(), - getStaticNumElements(op->getOperand(1).getType())}; - if (op->getNumOperands() > 0) - return {op->getOperand(0).getType(), elements}; - return {Type(), elements}; - }; - if (name == "tt.load" || name == "tt.store") { - bool load = name == "tt.load"; - auto [dataType, dataElements] = dataTypeAndElements(load); - int64_t bitWidth = dataType ? getTypeBitWidth(dataType) : 32; - double bytes = - static_cast(dataElements) * loopMultiplier * bitWidth / 8.0; - int64_t warpInstructions = - static_cast(std::ceil(dataElements / 32.0)) * loopMultiplier; - if (load) { - features.loadBytes += bytes; - features.loadWarpInstructions += warpInstructions; - if (inAnchor) { - features.simtAnchors.loadBytes += bytes; - features.simtAnchors.loadWarpInstructions += warpInstructions; - } - } else { - features.storeBytes += bytes; - features.storeWarpInstructions += warpInstructions; - if (inAnchor) { - features.simtAnchors.storeBytes += bytes; - features.simtAnchors.storeWarpInstructions += warpInstructions; - } - } - } - - if (name == "tt.dot" && op->getNumOperands() >= 2) { - auto lhs = dyn_cast(op->getOperand(0).getType()); - auto rhs = dyn_cast(op->getOperand(1).getType()); - if (lhs && rhs && lhs.getRank() >= 2 && rhs.getRank() >= 2) { - int64_t m = lhs.getShape()[lhs.getRank() - 2]; - int64_t k = lhs.getShape()[lhs.getRank() - 1]; - int64_t n = rhs.getShape()[rhs.getRank() - 1]; - if (m > 0 && n > 0 && k > 0) { - features.dotFlops += 2 * m * n * k * loopMultiplier; - features.dotOutputElements += m * n * loopMultiplier; - features.dotMNK.push_back({m, n, k}); - if (inAnchor) - features.simtAnchors.dotFlops += 2 * m * n * k * loopMultiplier; - } - } - } - - std::vector rankedResultAndOperandRanks; - bool hasRankedInput = false; - bool hasRankedResult = false; - for (Type type : op->getOperandTypes()) - if (auto tensor = dyn_cast(type)) { - rankedResultAndOperandRanks.push_back(tensor.getRank()); - hasRankedInput = true; - } - for (Type type : op->getResultTypes()) - if (auto tensor = dyn_cast(type)) { - rankedResultAndOperandRanks.push_back(tensor.getRank()); - hasRankedResult = true; - } - if (name == "tt.reduce") { - if (rankedResultAndOperandRanks.size() > 1) { - auto [minimum, maximum] = - std::minmax_element(rankedResultAndOperandRanks.begin(), - rankedResultAndOperandRanks.end()); - if (*maximum > *minimum) - ++features.rowLocalReduceOps; - } - if (hasRankedInput && !hasRankedResult) - ++features.vectorReduceToScalarOps; - if (op->getNumOperands() > 0) { - auto source = dyn_cast(op->getOperand(0).getType()); - auto axis = op->getAttrOfType("axis"); - if (source && source.hasStaticShape() && axis) { - int64_t axisValue = axis.getInt(); - if (axisValue < 0) - axisValue += source.getRank(); - if (axisValue >= 0 && axisValue < source.getRank()) { - int64_t extent = source.getShape()[axisValue]; - if (extent > 0) { - features.maxReduceAxisExtent = - std::max(features.maxReduceAxisExtent, extent); - features.weightedReduceAxisElements += extent * loopMultiplier; - const int64_t shuffleLevels = static_cast( - std::ceil(std::log2(static_cast(extent)))); - const int64_t inputElements = getStaticNumElements(source); - const int64_t shuffleLaneSteps = - inputElements * shuffleLevels * loopMultiplier; - features.shuffleLaneSteps += shuffleLaneSteps; - if (inAnchor) { - features.simtAnchors.maxReduceAxisExtent = - std::max(features.simtAnchors.maxReduceAxisExtent, extent); - features.simtAnchors.weightedReduceAxisElements += - extent * loopMultiplier; - features.simtAnchors.shuffleLaneSteps += shuffleLaneSteps; - } - } - } - } - } - } - - if ((name == "tt.scan" || name == "tt.associative_scan") && - op->getNumOperands() > 0) { - auto source = dyn_cast(op->getOperand(0).getType()); - auto axis = op->getAttrOfType("axis"); - if (source && source.hasStaticShape() && axis) { - int64_t axisValue = axis.getInt(); - if (axisValue < 0) - axisValue += source.getRank(); - if (axisValue >= 0 && axisValue < source.getRank()) { - int64_t extent = source.getShape()[axisValue]; - if (extent > 0) { - const int64_t shuffleLevels = static_cast( - std::ceil(std::log2(static_cast(extent)))); - const int64_t laneSteps = - getStaticNumElements(source) * shuffleLevels * loopMultiplier; - features.shuffleLaneSteps += laneSteps; - if (inAnchor) - features.simtAnchors.shuffleLaneSteps += laneSteps; - } - } - } - } - - std::vector maskRanks; - for (Type type : op->getOperandTypes()) - if (isMaskTensorType(type)) - maskRanks.push_back(cast(type).getRank()); - for (Type type : op->getResultTypes()) - if (isMaskTensorType(type)) - maskRanks.push_back(cast(type).getRank()); - if (!maskRanks.empty()) { - ++features.maskTensorOps; - for (int64_t rank : maskRanks) { - features.maskRankSum += rank; - if (inAnchor) - features.simtAnchors.maskRankSum += rank; - } - auto addPredicateLaneEvaluations = [&](Type type) { - if (!isMaskTensorType(type)) - return; - const int64_t laneEvaluations = - getStaticNumElements(type) * loopMultiplier; - features.predicateLaneEvaluations += laneEvaluations; - if (inAnchor) - features.simtAnchors.predicateLaneEvaluations += laneEvaluations; - }; - for (Type type : op->getOperandTypes()) - addPredicateLaneEvaluations(type); - for (Type type : op->getResultTypes()) - addPredicateLaneEvaluations(type); - if (name == "tt.broadcast" || name == "tt.expand_dims") - ++features.maskBroadcastOps; - } - - auto recordUniqueMask = [&](Value value) { - auto type = dyn_cast(value.getType()); - if (!type || !type.getElementType().isInteger(1)) - return; - if (uniqueMasks.insert(value).second) { - ++features.uniqueMaskValues; - features.uniqueMaskRankSum += type.getRank(); - const int64_t elements = getStaticNumElements(type); - features.predicateElements += elements; - } - if (inAnchor && anchorUniqueMasks.insert(value).second) { - ++features.simtAnchors.uniqueMaskValues; - features.simtAnchors.uniqueMaskRankSum += type.getRank(); - const int64_t elements = getStaticNumElements(type); - features.simtAnchors.predicateElements += elements; - } - }; - for (Value operand : op->getOperands()) - recordUniqueMask(operand); - for (Value result : op->getResults()) - recordUniqueMask(result); - - bool isPointerOperation = - name == "tt.addptr" || name == "tt.load" || name == "tt.store"; - if (isLoadedIndexDependentMemoryOp(op)) { + if (isLoadedIndexDependentMemoryOp(operation)) ++features.loadedIndexDependentMemoryOps; - if (inAnchor) - ++features.simtAnchors.loadedIndexDependentMemoryOps; - } - if (isPointerOperation) { - std::set uniqueShapes; - auto collectShape = [&](Type type) { - auto tensor = dyn_cast(type); - if (!tensor) - return; - std::string key; - llvm::raw_string_ostream os(key); - os << tensor.getRank(); - for (int64_t dim : tensor.getShape()) - os << 'x' << dim; - os.flush(); - uniqueShapes.insert(std::move(key)); - }; - for (Type type : op->getOperandTypes()) - collectShape(type); - for (Type type : op->getResultTypes()) - collectShape(type); - int64_t maxPointerRank = 0; - for (const std::string &shape : uniqueShapes) { - llvm::StringRef shapeRef(shape); - int64_t rank = 0; - (void)shapeRef.take_front(shapeRef.find('x')).getAsInteger(10, rank); - ++features.pointerTensorOps; - if (inAnchor) - ++features.simtAnchors.pointerTensorOps; - maxPointerRank = std::max(maxPointerRank, rank); - if (rank > 1) - features.pointerUnstructuredDims += rank; - } - if (maxPointerRank > 1) { - ++features.laneDependentPointerOps; - if (inAnchor) - ++features.simtAnchors.laneDependentPointerOps; + if (name == "tt.dot" && operation->getNumOperands() >= 2) { + auto lhs = dyn_cast(operation->getOperand(0).getType()); + auto rhs = dyn_cast(operation->getOperand(1).getType()); + if (lhs && rhs && lhs.getRank() >= 2 && rhs.getRank() >= 2) { + const int64_t m = lhs.getShape()[lhs.getRank() - 2]; + const int64_t k = lhs.getShape()[lhs.getRank() - 1]; + const int64_t n = rhs.getShape()[rhs.getRank() - 1]; + if (m > 0 && n > 0 && k > 0) + features.dotFlops += 2 * m * n * k * multiplier; } } - - bool anyRankedType = llvm::any_of(op->getOperandTypes(), [](Type type) { - return isa(type); - }); - anyRankedType |= llvm::any_of(op->getResultTypes(), [](Type type) { - return isa(type); - }); - if (name == "tt.load" && !anyRankedType && op->getNumOperands() > 0 && - isPointerType(op->getOperand(0).getType())) - ++features.scalarLoadOps; - if (name == "tt.store" && !anyRankedType && op->getNumOperands() > 0 && - isPointerType(op->getOperand(0).getType())) - ++features.scalarStoreOps; - if (name == "tt.splat" && op->getNumOperands() > 0 && - op->getNumResults() > 0 && isPointerType(op->getOperand(0).getType()) && - isa(op->getResult(0).getType())) - ++features.vectorPtrSplatOps; }); - - features.scalarOps = features.addOps + features.subOps + features.mulOps + - features.divOps + features.maxOps + features.absOps + - features.expOps + features.logOps + features.cmpOps + - features.selectOps + features.castOps + - features.clampOps; - features.hasDot = features.dotOps > 0; - features.hasGather = features.gatherOps > 0; - features.hasAtomic = features.atomicOps > 0; - features.hasHistogram = features.histogramOps > 0; - features.hasScan = features.scanOps > 0; - features.rank1IndirectVectorReduce = - features.maxTensorRank == 1 && features.reduceOps > 0 && - features.vectorReduceToScalarOps > 0 && features.vectorPtrSplatOps > 0 && - features.scalarLoadOps >= 2; - return features; } @@ -2402,8 +896,8 @@ estimateSimdSimtCandidatesImpl(const SimdSimtFeatureSummary &features, CandidateProfile profile = std::move(*profileOrError); SimdSimtCostReport report; - report.profileVersion = profile.profileVersion; - report.profileTarget = profile.target; + report.profileVersion = profile.hardware.profileVersion; + report.profileTarget = profile.hardware.target; report.actualTarget = options.actualTarget; report.profileContentSha256 = profile.contentSha256; report.selectionProfileContentSha256 = profile.selectionContentSha256; @@ -2412,41 +906,25 @@ estimateSimdSimtCandidatesImpl(const SimdSimtFeatureSummary &features, report.microbenchmarkProfileContentSha256 = profile.microbenchmarkContentSha256; report.scoreUnit = profile.scoreUnit; - report.targetCompatible = targetMatches(profile, options.actualTarget); report.features = features; - report.applicability = - evaluateSimtApplicability(features, options.compileOn91095); report.allSimdCandidateLegal = - features.simtAnchors.kernelLowerability.allSimd == - CandidateLoweringStatus::Native; - const CandidateLoweringStatus allSimtStatus = - features.simtAnchors.kernelLowerability.allSimtOnly; - // BackendConditional means the anchor itself has no TTIR legality blocker; - // the 910/95 pure-SIMT backend is the remaining condition. Once that - // backend is selected, let the complete Stage model decide the route. + features.simtAnchors.kernelLowerability.allSimd; report.allSimtOnlyCandidateLegal = options.compileOn91095 && !features.hasExplicitScope && - (allSimtStatus == CandidateLoweringStatus::Native || - allSimtStatus == CandidateLoweringStatus::BackendConditional); + features.simtAnchors.kernelLowerability.allSimtOnly; report.mixedCandidateLegal = !features.hasExplicitScope && - report.applicability.materializable && - features.simtAnchors.kernelLowerability.mixed == - CandidateLoweringStatus::Native; + options.compileOn91095 && + features.simtAnchors.count > 0 && + features.simtAnchors.kernelLowerability.mixed; report.includeFeaturesInJSON = options.includeFeaturesInJSON; - const int64_t weightedReductions = - mapValue(features.weightedOps, "reduce", features.reduceOps); - const int64_t dotFlops = features.dotFlops; - const int64_t pointerOps = std::max(1, features.pointerTensorOps); - report.breakdown.irregularDensity = std::min( - 1.0, static_cast(features.laneDependentPointerOps) / pointerOps); - const int64_t numWarps = std::max(1, static_cast(options.numWarps)); auto stageModel = evaluateStageModel( features, profile, static_cast(numWarps), options.wholeKernelSuperblockMaterializable, - options.scopeSuperblockMaterializable, module, anchorPlan); + options.scopeSuperblockMaterializable, options.logicalProgramCountHint, + options.physicalVectorCoreCountHint, module, anchorPlan); if (!stageModel) return stageModel.takeError(); if (*stageModel) { @@ -2457,7 +935,6 @@ estimateSimdSimtCandidatesImpl(const SimdSimtFeatureSummary &features, report.allSimdCandidateLegal &= report.stageModel.allSimd.legal; report.allSimtOnlyCandidateLegal &= report.stageModel.allSimt.legal; report.mixedCandidateLegal &= report.stageModel.mixed.legal; - report.breakdown.mixedCostSource = "stage_cost_evaluator_route_sum"; const unsigned legalCandidateCount = static_cast(report.allSimdCandidateLegal) + static_cast(report.allSimtOnlyCandidateLegal) + @@ -2469,421 +946,17 @@ estimateSimdSimtCandidatesImpl(const SimdSimtFeatureSummary &features, report.decision = chooseBest( report.candidateCosts, report.allSimdCandidateLegal, report.allSimtOnlyCandidateLegal, report.mixedCandidateLegal); - report.bestScore = report.candidateCosts.get(report.decision); - const double denominator = std::max(1.0e-9, report.bestScore); - report.candidateRatiosToBest = { - report.candidateCosts.allSimd / denominator, - report.candidateCosts.allSimtOnly / denominator, - report.candidateCosts.mixedSimdSimt / denominator}; return report; } - const int64_t maxNumel = std::max(1, features.maxTensorNumel); - const int64_t elementBits = - features.maxElementBits > 0 - ? std::max(8, features.maxElementBits) - : 32; - const int64_t vectorWidth = - std::max(1, profile.simdVectorWidthBits / elementBits); - - llvm::StringMap rawCountByKind; - rawCountByKind["gather"] = features.gatherOps; - rawCountByKind["histogram"] = features.histogramOps; - rawCountByKind["atomic"] = features.atomicOps; - for (llvm::StringRef kind : {"gather", "histogram", "atomic"}) { - int64_t coreWork = mapValue(features.opElements, kind, - mapValue(rawCountByKind, kind) * maxNumel); - if (coreWork > 0) - report.unsupported.push_back((kind + "_core_cost_uncalibrated").str()); - } - - int64_t classifiedScalarOps = - features.addOps + features.subOps + features.mulOps + features.divOps + - features.maxOps + features.absOps + features.expOps + features.logOps + - features.cmpOps + features.selectOps + features.castOps + - features.clampOps; - int64_t unclassifiedScalarOps = - std::max(0, features.scalarOps - classifiedScalarOps); - if (unclassifiedScalarOps) - report.unsupported.push_back(std::to_string(unclassifiedScalarOps) + - " unclassified arithmetic ops"); - - for (const auto &[opName, elements] : getProfileOpElements(features)) { - if (elements <= 0) - continue; - auto simdIterator = profile.simdOps.find(opName); - auto simtIterator = profile.simtOps.find(opName); - if (simdIterator == profile.simdOps.end() || - simtIterator == profile.simtOps.end()) { - report.unsupported.push_back(opName.str()); - continue; - } - const OpProfile &simd = simdIterator->second; - const OpProfile &simt = simtIterator->second; - if (simd.throughput <= 0.0 || simt.throughput <= 0.0) { - report.unsupported.push_back(opName.str()); - continue; - } - double simdCycles = std::ceil(static_cast(elements) / vectorWidth) / - simd.throughput * simd.factor; - double simtCycles = - static_cast(elements) / simt.throughput * simt.factor; - report.breakdown.simdOpSystemCycles[opName] = simdCycles; - report.breakdown.simtOpSystemCycles[opName] = simtCycles; - report.breakdown.simdComputeCycles += simdCycles; - report.breakdown.simtComputeCycles += simtCycles; - } - - for (const auto &[opName, elements] : - getProfileOpElements(features.simtAnchors)) { - if (elements <= 0) - continue; - auto simdIterator = profile.simdOps.find(opName); - auto simtIterator = profile.simtOps.find(opName); - if (simdIterator == profile.simdOps.end() || - simtIterator == profile.simtOps.end()) - continue; - const OpProfile &simd = simdIterator->second; - const OpProfile &simt = simtIterator->second; - if (simd.throughput <= 0.0 || simt.throughput <= 0.0) - continue; - report.breakdown.mixedSimdRegularComputeCycles -= - std::ceil(static_cast(elements) / vectorWidth) / - simd.throughput * simd.factor; - report.breakdown.mixedSimtAnchorComputeCycles += - static_cast(elements) / simt.throughput * simt.factor; - } - report.breakdown.mixedSimdRegularComputeCycles += - report.breakdown.simdComputeCycles; - report.breakdown.mixedSimdRegularComputeCycles = - std::max(0.0, report.breakdown.mixedSimdRegularComputeCycles); - - report.breakdown.simdLoadCycles = - features.loadBytes / profile.simdMte2BytesPerCycle; - report.breakdown.simdStoreCycles = - features.storeBytes / profile.simdMte3BytesPerCycle; - report.breakdown.simdMemoryCycles = std::max( - report.breakdown.simdLoadCycles, report.breakdown.simdStoreCycles); - const double mixedSimdRegularLoadCycles = - std::max(0.0, features.loadBytes - features.simtAnchors.loadBytes) / - profile.simdMte2BytesPerCycle; - const double mixedSimdRegularStoreCycles = - std::max(0.0, features.storeBytes - features.simtAnchors.storeBytes) / - profile.simdMte3BytesPerCycle; - report.breakdown.mixedSimdRegularMemoryCycles = - std::max(mixedSimdRegularLoadCycles, mixedSimdRegularStoreCycles); - - const int64_t loadWarpInstructions = - features.loadWarpInstructions != 0 - ? features.loadWarpInstructions - : features.loadOps * - static_cast(std::ceil(static_cast(maxNumel) / - profile.simtWarpSize)); - const int64_t storeWarpInstructions = - features.storeWarpInstructions != 0 - ? features.storeWarpInstructions - : features.storeOps * - static_cast(std::ceil(static_cast(maxNumel) / - profile.simtWarpSize)); - report.features.loadWarpInstructions = loadWarpInstructions; - report.features.storeWarpInstructions = storeWarpInstructions; - report.breakdown.simtLoadCycles = - loadWarpInstructions / profile.simtLoadWarpRate; - report.breakdown.simtStoreCycles = - storeWarpInstructions / profile.simtStoreWarpRate; - report.breakdown.simtMemoryCycles = - report.breakdown.simtLoadCycles + report.breakdown.simtStoreCycles; - report.breakdown.mixedSimtAnchorMemoryCycles = - features.simtAnchors.loadWarpInstructions / profile.simtLoadWarpRate + - features.simtAnchors.storeWarpInstructions / profile.simtStoreWarpRate; - - const int64_t weightedScans = - mapValue(features.weightedOps, "scan", features.scanOps); - if (weightedScans) - report.unsupported.push_back("scan_template_ranking_uncalibrated"); - const int64_t shuffleLevels = static_cast( - std::ceil(std::log2(static_cast(profile.simtWarpSize)))); - report.breakdown.simtShuffleInstructions = - features.shuffleLaneSteps > 0 - ? std::ceil(static_cast(features.shuffleLaneSteps) / - profile.simtWarpSize) - : static_cast(weightedReductions + weightedScans) * - std::ceil(static_cast(maxNumel) / - profile.simtWarpSize) * - shuffleLevels; - report.breakdown.simtShuffleCycles = - report.breakdown.simtShuffleInstructions / profile.simtShuffleRate; - const int64_t anchorWeightedReductions = - mapValue(features.simtAnchors.weightedOps, "reduce", - features.simtAnchors.reduceOps); - const int64_t anchorWeightedScans = mapValue( - features.simtAnchors.weightedOps, "scan", features.simtAnchors.scanOps); - const int64_t anchorMaxNumel = - std::max(1, features.simtAnchors.maxTensorNumel); - const double anchorShuffleInstructions = - features.simtAnchors.shuffleLaneSteps > 0 - ? std::ceil( - static_cast(features.simtAnchors.shuffleLaneSteps) / - profile.simtWarpSize) - : static_cast(anchorWeightedReductions + - anchorWeightedScans) * - std::ceil(static_cast(anchorMaxNumel) / - profile.simtWarpSize) * - shuffleLevels; - report.breakdown.mixedSimtAnchorShuffleCycles = - anchorShuffleInstructions / profile.simtShuffleRate; - - report.breakdown.simtPredicateInstructions = - features.predicateLaneEvaluations > 0 - ? std::ceil(static_cast(features.predicateLaneEvaluations) / - profile.simtWarpSize) - : static_cast(features.maskRankSum) * - std::ceil(static_cast(maxNumel) / profile.simtWarpSize); - report.breakdown.simtPredicateCycles = - report.breakdown.simtPredicateInstructions / profile.simtPredicateRate; - const double anchorPredicateInstructions = - features.simtAnchors.predicateLaneEvaluations > 0 - ? std::ceil(static_cast( - features.simtAnchors.predicateLaneEvaluations) / - profile.simtWarpSize) - : static_cast(features.simtAnchors.maskRankSum) * - std::ceil(static_cast(anchorMaxNumel) / - profile.simtWarpSize); - report.breakdown.mixedSimtAnchorPredicateCycles = - anchorPredicateInstructions / profile.simtPredicateRate; - - if (dotFlops) { - report.breakdown.simdDotCycles = - profile.simdDotSetupCycles + - static_cast(dotFlops) / profile.simdDotFlopsPerCycle; - report.breakdown.simtDotCycles = - profile.simtDotSetupCycles + - static_cast(dotFlops) / profile.simtDotFlopsPerCycle; - } - const int64_t regularDotFlops = - std::max(0, dotFlops - features.simtAnchors.dotFlops); - report.breakdown.cubeTailDotOps = - std::max(0, features.dotOps - features.simtAnchors.dotOps); - report.breakdown.cubeTailDotFlops = regularDotFlops; - if (regularDotFlops) - report.breakdown.mixedSimdRegularDotCycles = - profile.simdDotSetupCycles + - static_cast(regularDotFlops) / profile.simdDotFlopsPerCycle; - if (features.simtAnchors.dotFlops) - report.breakdown.mixedSimtAnchorDotCycles = - profile.simtDotSetupCycles + - static_cast(features.simtAnchors.dotFlops) / - profile.simtDotFlopsPerCycle; - - report.breakdown.simdSetupCycles = profile.simdSetupCycles; - report.breakdown.simtSetupCycles = profile.simtSetupCycles; - report.breakdown.simdIssuePayloadCycles = std::max( - report.breakdown.simdComputeCycles + report.breakdown.simdDotCycles, - report.breakdown.simdMemoryCycles); - // The current SIMT lowering emits a dependency-ordered warp instruction - // stream. Loads feed compute, compute feeds shuffle/reduction, and stores - // consume the result. There is no measured overlap contract that would - // justify a roofline max(compute, memory), so charge the serial path. - report.breakdown.simtIssuePayloadCycles = - report.breakdown.simtComputeCycles + report.breakdown.simtShuffleCycles + - report.breakdown.simtDotCycles + report.breakdown.simtMemoryCycles + - report.breakdown.simtPredicateCycles; - report.breakdown.programIssueScale = profile.programIssueScale; - report.breakdown.simdAnalyticalCycles = - profile.simdSetupCycles + - report.breakdown.simdIssuePayloadCycles * profile.programIssueScale; - report.breakdown.simtAnalyticalCycles = - profile.simtSetupCycles + - report.breakdown.simtIssuePayloadCycles * profile.programIssueScale; - - const bool tinyDot = - dotFlops > 0 && dotFlops <= profile.structural.tinyDotFlopsMax; - report.breakdown.tinyDotUnderfill = - tinyDot ? std::max(0.0, 1.0 - static_cast(dotFlops) / - profile.structural.tinyDotFlopsMax) - : 0.0; - const double irregularPerDensity = - tinyDot ? profile.structural.tinyDotIrregularPerDensity - : profile.structural.irregularPerDensity; - const double irregularCap = tinyDot ? profile.structural.tinyDotIrregularCap - : profile.structural.irregularCap; - report.breakdown.structuralComponents["irregular_addressing"] = std::min( - irregularCap, report.breakdown.irregularDensity * irregularPerDensity); - report.breakdown.structuralComponents["mask_materialization"] = - std::min(profile.structural.maskCap, - features.maskRankSum * profile.structural.perMaskRank); - report.breakdown.structuralComponents["reduction_lowering"] = - std::min(profile.structural.reductionCap, - weightedReductions * profile.structural.perWeightedReduction); - report.breakdown.structuralComponents["static_loop_control"] = std::min( - profile.structural.loopCap, - features.staticLoopTripCountSum * profile.structural.perStaticLoopTrip); - report.breakdown.structuralComponents["control_flow"] = - features.hasControlFlow ? profile.structural.controlFlow : 0.0; - report.breakdown.structuralComponents["tiny_dot_startup"] = - tinyDot ? profile.structural.tinyDot * report.breakdown.tinyDotUnderfill - : 0.0; - report.breakdown.structuralComponents["rank1_indirect_vector_reduction"] = - features.rank1IndirectVectorReduce - ? profile.structural.rank1IndirectVectorReduction - : 0.0; - for (const auto &component : report.breakdown.structuralComponents) - report.breakdown.structuralPenaltyRatio += component.second; - // Candidate costs must remain independent. Structural terms describe work - // omitted by the SIMD roofline, so charge them against A_SIMD itself; - // changing SIMT throughput/setup must never change the all-SIMD score. - report.breakdown.simdStructuralPenaltyCycles = - report.breakdown.simdAnalyticalCycles * - report.breakdown.structuralPenaltyRatio; - report.candidateCosts.allSimd = report.breakdown.simdAnalyticalCycles + - report.breakdown.simdStructuralPenaltyCycles; - report.candidateCosts.allSimtOnly = report.breakdown.simtAnalyticalCycles; - const MixedSetupFallbackProfile *nearestSetupFallback = nullptr; - for (const MixedSetupFallbackProfile &fallback : profile.mixedSetupFallbacks) - if (!nearestSetupFallback || - std::abs(fallback.numWarps - numWarps) < - std::abs(nearestSetupFallback->numWarps - numWarps)) - nearestSetupFallback = &fallback; - if (!nearestSetupFallback) - return llvm::createStringError(std::errc::invalid_argument, - "SIMD/SIMT profile has no mixed setup " - "fallback"); - report.breakdown.mixedSetupFallbackNumWarps = nearestSetupFallback->numWarps; - report.breakdown.standaloneSimtSetupCycles = profile.simtSetupCycles; - report.breakdown.mixedSetupFallbackCycles = - nearestSetupFallback->emptySimtSetupCycles; - report.breakdown.setupProxyDeltaCycles = - std::max(0.0, report.breakdown.mixedSetupFallbackCycles - - report.breakdown.standaloneSimtSetupCycles); - - // A mixed route is not a convex blend of two whole-kernel costs. Charge - // exact materializable anchors at SIMT rates and remaining operations at - // SIMD rates. The regular SIMD phase keeps its measured roofline model; - // the SIMT anchor is a serial load/compute/shuffle/store instruction path. - report.breakdown.mixedSimdRegularPayloadCycles = - std::max(report.breakdown.mixedSimdRegularComputeCycles + - report.breakdown.mixedSimdRegularDotCycles, - report.breakdown.mixedSimdRegularMemoryCycles); - report.breakdown.mixedSimtAnchorPayloadCycles = - report.breakdown.mixedSimtAnchorComputeCycles + - report.breakdown.mixedSimtAnchorDotCycles + - report.breakdown.mixedSimtAnchorShuffleCycles + - report.breakdown.mixedSimtAnchorMemoryCycles + - report.breakdown.mixedSimtAnchorPredicateCycles; - - const int64_t remainingPointerOps = std::max( - 0, features.pointerTensorOps - features.simtAnchors.pointerTensorOps); - const int64_t remainingLaneDependentPointerOps = - std::max(0, features.laneDependentPointerOps - - features.simtAnchors.laneDependentPointerOps); - const double remainingIrregularDensity = - remainingPointerOps > 0 - ? std::min(1.0, - static_cast(remainingLaneDependentPointerOps) / - remainingPointerOps) - : 0.0; - const int64_t remainingMaskRank = std::max( - 0, features.maskRankSum - features.simtAnchors.maskRankSum); - const int64_t remainingWeightedReductions = - std::max(0, weightedReductions - anchorWeightedReductions); - const int64_t remainingLoopTrips = - std::max(0, features.staticLoopTripCountSum - - features.simtAnchors.staticLoopTripCountSum); - const bool remainingControlFlow = - features.hasControlFlow && !features.simtAnchors.hasControlFlow; - const bool remainingRank1Reduction = - features.rank1IndirectVectorReduce && remainingWeightedReductions > 0; - const bool remainingTinyDot = regularDotFlops > 0 && tinyDot; - - double remainingStructuralPenalty = 0.0; - remainingStructuralPenalty += - std::min(irregularCap, remainingIrregularDensity * irregularPerDensity); - remainingStructuralPenalty += - std::min(profile.structural.maskCap, - remainingMaskRank * profile.structural.perMaskRank); - remainingStructuralPenalty += std::min( - profile.structural.reductionCap, - remainingWeightedReductions * profile.structural.perWeightedReduction); - remainingStructuralPenalty += - std::min(profile.structural.loopCap, - remainingLoopTrips * profile.structural.perStaticLoopTrip); - if (remainingControlFlow) - remainingStructuralPenalty += profile.structural.controlFlow; - if (remainingRank1Reduction) - remainingStructuralPenalty += - profile.structural.rank1IndirectVectorReduction; - if (remainingTinyDot) - remainingStructuralPenalty += - profile.structural.tinyDot * report.breakdown.tinyDotUnderfill; - report.breakdown.mixedRemainingStructuralPenaltyRatio = - remainingStructuralPenalty; - - double totalPartitionWork = features.loadBytes + features.storeBytes + - static_cast(features.dotFlops); - for (const auto &entry : features.opElements) - totalPartitionWork += std::max(0, entry.second); - double anchorPartitionWork = - features.simtAnchors.loadBytes + features.simtAnchors.storeBytes + - static_cast(features.simtAnchors.dotFlops); - for (const auto &entry : features.simtAnchors.opElements) - anchorPartitionWork += std::max(0, entry.second); - report.breakdown.mixedSimdFraction = - totalPartitionWork > 0.0 - ? std::clamp(1.0 - anchorPartitionWork / totalPartitionWork, 0.0, 1.0) - : 0.0; - - if (features.simtAnchors.count > 0) { - const double regularPayloadWithResidual = - report.breakdown.mixedSimdRegularPayloadCycles * - (1.0 + remainingStructuralPenalty); - double anchorPayloadForCost = report.breakdown.mixedSimtAnchorPayloadCycles; - report.breakdown.mixedSimtAnchorCalibratedPayloadCycles = - anchorPayloadForCost; - report.breakdown.mixedSimtAnchorCalibratedPayloadCycles = - anchorPayloadForCost; - report.candidateCosts.mixedSimdSimt = - report.breakdown.mixedSetupFallbackCycles + - profile.programIssueScale * - (regularPayloadWithResidual + anchorPayloadForCost) + - report.breakdown.mixedBoundaryCycles; - report.breakdown.mixedCostSource = - "materializable_anchor_resource_partition"; - } else { - report.candidateCosts.mixedSimdSimt = - std::max(report.candidateCosts.allSimd, - report.candidateCosts.allSimtOnly) + - report.breakdown.mixedSetupFallbackCycles; - report.breakdown.mixedCostSource = - "inapplicable_without_materializable_anchor"; - } - - sortAndUnique(report.unsupported); - const unsigned legalCandidateCount = - static_cast(report.allSimdCandidateLegal) + - static_cast(report.allSimtOnlyCandidateLegal) + - static_cast(report.mixedCandidateLegal); - if (legalCandidateCount == 0) - return llvm::createStringError( - std::errc::not_supported, - "SIMD/SIMT route model found no independently lowerable candidate"); - report.decision = - chooseBest(report.candidateCosts, report.allSimdCandidateLegal, - report.allSimtOnlyCandidateLegal, report.mixedCandidateLegal); - report.bestScore = report.candidateCosts.get(report.decision); - const double ratioDenominator = std::max(1.0e-9, report.bestScore); - report.candidateRatiosToBest = { - report.candidateCosts.allSimd / ratioDenominator, - report.candidateCosts.allSimtOnly / ratioDenominator, - report.candidateCosts.mixedSimdSimt / ratioDenominator}; + // Unknown Stage domains are deliberately not scored. The online Route + // Model has no legacy whole-kernel fallback; the selector leaves the + // existing backend-default lowering unchanged. + report.stageModel.applied = false; + report.unsupported.push_back("stage_model_not_applicable"); return report; } -llvm::Expected mlir::ascend::estimateSimdSimtCandidates( - const SimdSimtFeatureSummary &features, - const SimdSimtCostModelOptions &options) { - return estimateSimdSimtCandidatesImpl(features, options, nullptr, nullptr); -} - llvm::Expected mlir::ascend::analyzeSimdSimtCandidates( ModuleOp module, const SimdSimtCostModelOptions &options) { if (!module) diff --git a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageCostModels.cpp b/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageCostModels.cpp index 1d9c0dac42..ed5f7a7966 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageCostModels.cpp +++ b/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageCostModels.cpp @@ -4,7 +4,6 @@ #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/StringSet.h" -#include "llvm/ADT/StringSwitch.h" #include "llvm/Support/ErrorHandling.h" #include @@ -23,7 +22,7 @@ static double iterations(const LogicalStage &stage) { } static double controlBody(const StageResourceCycles &resources) { - return resources.control + resources.loopControl + resources.branchControl + + return resources.loopControl + resources.branchControl + resources.divergence + resources.synchronization; } @@ -37,11 +36,6 @@ static double serialBody(const StageResourceCycles &resources) { return std::max(execution, resources.issue); } -static bool supportsAny(StageCostModelKind kind, - std::initializer_list kinds) { - return llvm::is_contained(kinds, kind); -} - static bool permitsSimdOverlap(const LogicalStage &stage) { return stage.scheduleKind == StageScheduleKind::IndependentPipelined && stage.features.permitsSimdRoofline(); @@ -69,83 +63,36 @@ materializeControlFlow(const LogicalStage &stage, StageMode mode, return resources; } -static StageResourceCycles mapSIMDWorkload(const LogicalStage &stage, - const StageModeProfile &profile) { +static StageResourceCycles mapWorkload(const LogicalStage &stage, + const StageModeProfile &profile, + StageMode mode) { StageResourceCycles resources; const StageWorkload &work = stage.workload; + const bool simd = mode == StageMode::SIMD; resources.setup = work.paysKernelSetup ? profile.setupCycles : 0.0; for (const auto &[name, elements] : work.operationElements) { auto rate = profile.operationRates.find(name); if (rate == profile.operationRates.end() || rate->second.throughput <= 0.0) continue; + const double instructions = + simd ? std::ceil(elements / static_cast(profile.vectorWidth)) + : elements; resources.compute += - std::ceil(elements / static_cast(profile.vectorWidth)) / - rate->second.throughput * rate->second.factor; + instructions / rate->second.throughput * rate->second.factor; } resources.scalar = work.scalarOperations / profile.scalarOperationsPerCycle; if (stage.features.hasIndirectMemory) { - const double loadTransactions = + const double loads = std::max(work.loadWarpInstructions, work.loadBytes > 0.0 ? 1.0 : 0.0); - const double storeTransactions = + const double stores = std::max(work.storeWarpInstructions, work.storeBytes > 0.0 ? 1.0 : 0.0); - resources.load = - loadTransactions / profile.indirectLoadTransactionsPerCycle; - resources.store = - storeTransactions / profile.indirectStoreTransactionsPerCycle; - if (loadTransactions + storeTransactions > 0.0) + resources.load = loads / profile.indirectLoadTransactionsPerCycle; + resources.store = stores / profile.indirectStoreTransactionsPerCycle; + if (loads + stores > 0.0) resources.load += profile.indirectDependencyLatencyCycles; - } else { + } else if (simd) { resources.load = work.loadBytes / profile.loadBytesPerCycle; resources.store = work.storeBytes / profile.storeBytesPerCycle; - } - resources.predicate = std::ceil(work.predicateElements / - static_cast(profile.vectorWidth)) / - profile.predicateOperationsPerCycle; - resources.shuffle = work.shuffleLaneSteps / profile.shuffleLanesPerCycle; - if (work.dotFlops > 0.0) { - resources.setup += profile.dotSetupCycles; - resources.dot = work.dotFlops / profile.dotFlopsPerCycle; - } - resources.issue = - std::ceil(work.issueElements / static_cast(profile.issueWidth)) / - profile.issueOperationsPerCycle; - resources.spill = - work.estimatedSpillTransactions / profile.spillTransactionsPerCycle; - if (stage.features.hasLoopCarriedDataDependency) - resources.criticalPath = resources.scalar + resources.compute + - resources.predicate + resources.shuffle + - resources.dot; - else if (stage.features.hasReduction) - resources.criticalPath = - resources.compute + resources.predicate + resources.shuffle; - return materializeControlFlow(stage, StageMode::SIMD, resources, - profile.controlFlow); -} - -static StageResourceCycles mapSIMTWorkload(const LogicalStage &stage, - const StageModeProfile &profile) { - StageResourceCycles resources; - const StageWorkload &work = stage.workload; - resources.setup = work.paysKernelSetup ? profile.setupCycles : 0.0; - for (const auto &[name, elements] : work.operationElements) { - auto rate = profile.operationRates.find(name); - if (rate == profile.operationRates.end() || rate->second.throughput <= 0.0) - continue; - resources.compute += - elements / rate->second.throughput * rate->second.factor; - } - resources.scalar = work.scalarOperations / profile.scalarOperationsPerCycle; - if (stage.features.hasIndirectMemory) { - const double loadTransactions = - std::max(work.loadWarpInstructions, work.loadBytes > 0.0 ? 1.0 : 0.0); - const double storeTransactions = - std::max(work.storeWarpInstructions, work.storeBytes > 0.0 ? 1.0 : 0.0); - resources.load = - loadTransactions / profile.indirectLoadTransactionsPerCycle; - resources.store = - storeTransactions / profile.indirectStoreTransactionsPerCycle; - if (loadTransactions + storeTransactions > 0.0) - resources.load += profile.indirectDependencyLatencyCycles; } else { resources.load = work.loadWarpInstructions / profile.loadWarpInstructionsPerCycle; @@ -153,7 +100,10 @@ static StageResourceCycles mapSIMTWorkload(const LogicalStage &stage, work.storeWarpInstructions / profile.storeWarpInstructionsPerCycle; } resources.predicate = - work.predicateElements / profile.predicateOperationsPerCycle; + (simd ? std::ceil(work.predicateElements / + static_cast(profile.vectorWidth)) + : work.predicateElements) / + profile.predicateOperationsPerCycle; resources.shuffle = work.shuffleLaneSteps / profile.shuffleLanesPerCycle; if (work.dotFlops > 0.0) { resources.setup += profile.dotSetupCycles; @@ -171,8 +121,7 @@ static StageResourceCycles mapSIMTWorkload(const LogicalStage &stage, else if (stage.features.hasReduction) resources.criticalPath = resources.compute + resources.predicate + resources.shuffle; - return materializeControlFlow(stage, StageMode::SIMT, resources, - profile.controlFlow); + return materializeControlFlow(stage, mode, resources, profile.controlFlow); } static double applySuperBlock(const LogicalStage &stage, @@ -183,6 +132,7 @@ static double applySuperBlock(const LogicalStage &stage, if (implementation.mode != StageMode::SIMT || implementation.superblockFactor == 1) return stageCycles; + const double factor = static_cast(implementation.superblockFactor); const double effectiveFactor = std::min( factor, static_cast(profile.superblockUsefulFactorLimit)); @@ -196,12 +146,13 @@ static double applySuperBlock(const LogicalStage &stage, // divide dependent arithmetic, loop control, or synchronization. const double pressure = iterations(stage) * resources.spill * std::max(0.0, factor - 1.0); - // A loop-carried Stage keeps its live-out state across iterations. Factors - // above the target's pressure-free point replicate that persistent state - // across more logical warp groups. This is a resource cost, not a blanket - // per-factor penalty: straight-line and small-state Stages are unaffected. + // Live-out bytes alone do not prove register pressure: they describe the + // Stage ABI, not the allocator's simultaneously-live set. Charge replicated + // persistent state only when workload analysis has independently predicted + // spill traffic. This keeps the penalty evidence based and lets independent + // recurrence groups use F4 when the generated SIMT VF has no STK/LDK. const double persistentStatePressure = - stage.features.hasLoopCarriedDataDependency + stage.features.hasLoopCarriedDataDependency && resources.spill > 0.0 ? std::max( 0.0, factor - @@ -210,358 +161,111 @@ static double applySuperBlock(const LogicalStage &stage, static_cast(stage.liveOutBytes) / profile.superblockPersistentStateBytesPerCycle : 0.0; - const double issueFloor = resources.setup + - iterations(stage) * resources.issue + - resources.epilogue; - // Persistent-state pressure is additional register/stack work and cannot - // disappear behind the ordinary issue floor. - return std::max(issueFloor, stageCycles - latencySensitive + - latencySensitive / effectiveFactor + - pressure) + + const double fixed = resources.setup; + const double issueFloor = + fixed + factor * iterations(stage) * resources.issue; + // A recurrence is serial inside one logical program. SuperBlock contributes + // F independent logical programs to the same physical program, allowing the + // scheduler to cover one program's dependency stalls with another program. + // Normalize the critical-path portion per logical program, but retain the + // aggregate issue floor: F2/F4 cannot create additional issue bandwidth. + // This applies equally to whole-kernel and scope-local SuperBlock because + // both materializers batch complete logical programs around the Stage. + if (stage.costModelKind == StageCostModelKind::LoopCarriedRecurrence) { + const double recurrenceBody = std::max(0.0, stageCycles - fixed); + return std::max(issueFloor, fixed + recurrenceBody + pressure) + + persistentStatePressure; + } + // Proven persistent-state pressure is additional register/stack work and + // cannot disappear behind the ordinary issue floor. + const double body = std::max(0.0, stageCycles - fixed); + const double groupedBody = factor * std::max(0.0, body - latencySensitive) + + factor * latencySensitive / effectiveFactor; + return std::max(issueFloor, fixed + groupedBody + pressure) + persistentStatePressure; } -class SIMDDispatchStageCostModel final : public SIMDStageCostModel { -public: - llvm::StringRef getName() const override { return "simd_dispatch"; } - bool supports(StageCostModelKind kind) const override { - return supportsAny(kind, {StageCostModelKind::AutoBlockifyDispatch, - StageCostModelKind::AutoBlockifyLoop}); - } - double estimate(const StageCostModelContext &context, - const StageImplementation &implementation, - const StageResourceCycles &resources) const override { - const StageResourceCycles &r = resources; - const double count = - context.stage.costModelKind == StageCostModelKind::AutoBlockifyLoop - ? iterations(context.stage) - : 1.0; - return r.setup + count * std::max(r.scalar + controlBody(r), r.issue) + - r.epilogue; - } -}; - -class SIMTDispatchStageCostModel final : public SIMTStageCostModel { -public: - llvm::StringRef getName() const override { return "simt_dispatch"; } - bool supports(StageCostModelKind kind) const override { - return supportsAny(kind, {StageCostModelKind::AutoBlockifyDispatch, - StageCostModelKind::AutoBlockifyLoop}); - } - double estimate(const StageCostModelContext &context, - const StageImplementation &implementation, - const StageResourceCycles &resources) const override { - const StageResourceCycles &r = resources; - const double count = - context.stage.costModelKind == StageCostModelKind::AutoBlockifyLoop - ? iterations(context.stage) - : 1.0; - return r.setup + count * std::max(r.scalar + controlBody(r), r.issue) + - r.epilogue; - } -}; - -class SIMDScalarStageCostModel final : public SIMDStageCostModel { -public: - llvm::StringRef getName() const override { return "simd_scalar"; } - bool supports(StageCostModelKind kind) const override { - return supportsAny( - kind, - {StageCostModelKind::ScalarIssue, StageCostModelKind::ScalarControl, - StageCostModelKind::ScalarMath, StageCostModelKind::IndexGeneration, - StageCostModelKind::PredicateMask, StageCostModelKind::LoopPredicate}); - } - double estimate(const StageCostModelContext &context, - const StageImplementation &implementation, - const StageResourceCycles &resources) const override { - const StageResourceCycles &r = resources; - return r.setup + iterations(context.stage) * serialBody(r) + r.epilogue; - } -}; - -class SIMTScalarStageCostModel final : public SIMTStageCostModel { -public: - llvm::StringRef getName() const override { return "simt_scalar"; } - bool supports(StageCostModelKind kind) const override { - return supportsAny( - kind, - {StageCostModelKind::ScalarIssue, StageCostModelKind::ScalarControl, - StageCostModelKind::ScalarMath, StageCostModelKind::IndexGeneration, - StageCostModelKind::PredicateMask, StageCostModelKind::LoopPredicate}); - } - double estimate(const StageCostModelContext &context, - const StageImplementation &implementation, - const StageResourceCycles &resources) const override { - const StageResourceCycles &r = resources; - return r.setup + iterations(context.stage) * serialBody(r) + r.epilogue; - } -}; - -class SIMDContinuousMemoryStageCostModel final : public SIMDStageCostModel { -public: - llvm::StringRef getName() const override { return "simd_continuous_memory"; } - bool supports(StageCostModelKind kind) const override { - return supportsAny(kind, {StageCostModelKind::ContinuousTileMemory, - StageCostModelKind::ContinuousTileStore, - StageCostModelKind::ContinuousShortLoad, - StageCostModelKind::CachePolicyStore}); - } - double estimate(const StageCostModelContext &context, - const StageImplementation &implementation, - const StageResourceCycles &resources) const override { - const StageResourceCycles &r = resources; - double body = serialBody(r); - if (permitsSimdOverlap(context.stage)) - body = r.scalar + r.predicate + controlBody(r) + r.spill + - std::max({r.load, r.store, r.issue}); - return r.setup + iterations(context.stage) * body + r.epilogue; - } -}; - -class SIMTContinuousMemoryStageCostModel final : public SIMTStageCostModel { -public: - llvm::StringRef getName() const override { return "simt_continuous_memory"; } - bool supports(StageCostModelKind kind) const override { - return supportsAny(kind, {StageCostModelKind::ContinuousTileMemory, - StageCostModelKind::ContinuousTileStore, - StageCostModelKind::ContinuousShortLoad, - StageCostModelKind::CachePolicyStore}); - } - double estimate(const StageCostModelContext &context, - const StageImplementation &implementation, - const StageResourceCycles &resources) const override { - const StageResourceCycles &r = resources; - return r.setup + iterations(context.stage) * serialBody(r) + r.epilogue; - } -}; - -class SIMDIndirectMemoryStageCostModel final : public SIMDStageCostModel { -public: - llvm::StringRef getName() const override { return "simd_indirect_memory"; } - bool supports(StageCostModelKind kind) const override { - return supportsAny(kind, {StageCostModelKind::IndirectScalarMemory, - StageCostModelKind::IndirectGatherMemory}); - } - double estimate(const StageCostModelContext &context, - const StageImplementation &implementation, - const StageResourceCycles &resources) const override { - const StageResourceCycles &r = resources; - return r.setup + iterations(context.stage) * serialBody(r) + r.epilogue; - } -}; - -class SIMTIndirectMemoryStageCostModel final : public SIMTStageCostModel { -public: - llvm::StringRef getName() const override { return "simt_indirect_memory"; } - bool supports(StageCostModelKind kind) const override { - return supportsAny(kind, {StageCostModelKind::IndirectScalarMemory, - StageCostModelKind::IndirectGatherMemory}); - } - double estimate(const StageCostModelContext &context, - const StageImplementation &implementation, - const StageResourceCycles &resources) const override { - const StageResourceCycles &r = resources; - return r.setup + iterations(context.stage) * serialBody(r) + r.epilogue; - } -}; - -class SIMDIndependentStageCostModel final : public SIMDStageCostModel { -public: - llvm::StringRef getName() const override { - return "simd_independent_pipeline"; - } - bool supports(StageCostModelKind kind) const override { - return kind == StageCostModelKind::IndependentPipelinedLoop; - } - double estimate(const StageCostModelContext &context, - const StageImplementation &implementation, - const StageResourceCycles &resources) const override { - const StageResourceCycles &r = resources; - double body = serialBody(r); - if (permitsSimdOverlap(context.stage)) - body = std::max({r.load, r.store, r.compute + r.dot + r.shuffle, - r.scalar + r.predicate + controlBody(r), r.issue}) + - r.spill; - return r.setup + iterations(context.stage) * body + r.epilogue; - } -}; - -class SIMTIndependentStageCostModel final : public SIMTStageCostModel { -public: - llvm::StringRef getName() const override { - return "simt_independent_pipeline"; - } - bool supports(StageCostModelKind kind) const override { - return kind == StageCostModelKind::IndependentPipelinedLoop; - } - double estimate(const StageCostModelContext &context, - const StageImplementation &implementation, - const StageResourceCycles &resources) const override { - const StageResourceCycles &r = resources; - return r.setup + iterations(context.stage) * serialBody(r) + r.epilogue; - } -}; - -class SIMDRecurrenceStageCostModel final : public SIMDStageCostModel { -public: - llvm::StringRef getName() const override { return "simd_recurrence"; } - bool supports(StageCostModelKind kind) const override { - return kind == StageCostModelKind::LoopCarriedRecurrence; +static double estimateStage(const LogicalStage &stage, + const HardwareProfile &profile, StageMode mode, + const StageResourceCycles &r) { + const double count = iterations(stage); + const double serial = r.setup + count * serialBody(r); + switch (stage.costModelKind) { + case StageCostModelKind::AutoBlockifyDispatch: + case StageCostModelKind::AutoBlockifyLoop: { + const double dispatchCount = + stage.costModelKind == StageCostModelKind::AutoBlockifyLoop ? count + : 1.0; + return r.setup + + dispatchCount * std::max(r.scalar + controlBody(r), r.issue); } - double estimate(const StageCostModelContext &context, - const StageImplementation &implementation, - const StageResourceCycles &resources) const override { - const StageResourceCycles &r = resources; + case StageCostModelKind::ContinuousTileMemory: + case StageCostModelKind::ContinuousTileStore: + case StageCostModelKind::ContinuousShortLoad: + case StageCostModelKind::CachePolicyStore: + if (mode == StageMode::SIMD && permitsSimdOverlap(stage)) + return r.setup + count * (r.scalar + r.predicate + controlBody(r) + + r.spill + std::max({r.load, r.store, r.issue})); + return serial; + case StageCostModelKind::IndependentPipelinedLoop: + if (mode == StageMode::SIMD && permitsSimdOverlap(stage)) + return r.setup + + count * + (std::max({r.load, r.store, r.compute + r.dot + r.shuffle, + r.scalar + r.predicate + controlBody(r), r.issue}) + + r.spill); + return serial; + case StageCostModelKind::LoopCarriedRecurrence: { const double critical = r.criticalPath > 0.0 ? std::max(r.criticalPath + r.load + r.store + controlBody(r) + r.spill, r.issue) : serialBody(r); - return r.setup + iterations(context.stage) * critical + r.epilogue; - } -}; - -class SIMTRecurrenceStageCostModel final : public SIMTStageCostModel { -public: - llvm::StringRef getName() const override { return "simt_recurrence"; } - bool supports(StageCostModelKind kind) const override { - return kind == StageCostModelKind::LoopCarriedRecurrence; - } - double estimate(const StageCostModelContext &context, - const StageImplementation &implementation, - const StageResourceCycles &resources) const override { - const StageResourceCycles &r = resources; - const double criticalPerIteration = - r.criticalPath > 0.0 ? std::max(r.criticalPath + r.load + r.store + - controlBody(r) + r.spill, - r.issue) - : serialBody(r); - const int64_t independentGroups = std::max( - 1, context.stage.features.parallelRecurrenceGroupCount); - const int64_t interleavedGroups = std::max( - 1, std::min(independentGroups, context.profile.logicalWarpGroupCount)); - // A carried dependency serializes iterations within one group, not - // sibling recurrences with disjoint state. SIMT can interleave those - // groups, while the aggregate front-end issue stream remains a hard - // throughput floor over all iterations. - const double criticalIterations = std::ceil( - iterations(context.stage) / static_cast(interleavedGroups)); - const double criticalPath = criticalIterations * criticalPerIteration; - const double issueFloor = iterations(context.stage) * r.issue; - return r.setup + std::max(criticalPath, issueFloor) + r.epilogue; - } -}; - -class SIMDReductionStageCostModel final : public SIMDStageCostModel { -public: - llvm::StringRef getName() const override { return "simd_reduction"; } - bool supports(StageCostModelKind kind) const override { - return kind == StageCostModelKind::RowwiseReduction; - } - double estimate(const StageCostModelContext &context, - const StageImplementation &implementation, - const StageResourceCycles &resources) const override { - const StageResourceCycles &r = resources; - // Tree/shuffle depth is a dependency chain, while the shared issue rate - // remains a lower bound over the complete instruction stream. - const double execution = - r.scalar + r.load + r.store + r.criticalPath + controlBody(r) + r.spill; - return r.setup + iterations(context.stage) * std::max(execution, r.issue) + - r.epilogue; - } -}; - -class SIMTReductionStageCostModel final : public SIMTStageCostModel { -public: - llvm::StringRef getName() const override { return "simt_reduction"; } - bool supports(StageCostModelKind kind) const override { - return kind == StageCostModelKind::RowwiseReduction; - } - double estimate(const StageCostModelContext &context, - const StageImplementation &implementation, - const StageResourceCycles &resources) const override { - const StageResourceCycles &r = resources; - const double execution = - r.scalar + r.load + r.store + r.criticalPath + controlBody(r) + r.spill; - return r.setup + iterations(context.stage) * std::max(execution, r.issue) + - r.epilogue; - } -}; - -class SIMDCubeStageCostModel final : public SIMDStageCostModel { -public: - llvm::StringRef getName() const override { return "simd_cube"; } - bool supports(StageCostModelKind kind) const override { - return supportsAny(kind, {StageCostModelKind::CubeRoofline, - StageCostModelKind::TinyCubeRoofline}); - } - double estimate(const StageCostModelContext &context, - const StageImplementation &implementation, - const StageResourceCycles &resources) const override { - const StageResourceCycles &r = resources; - double body = serialBody(r); - if (permitsSimdOverlap(context.stage)) - body = r.scalar + r.predicate + controlBody(r) + r.shuffle + r.spill + - std::max({r.load, r.compute + r.dot, r.store, r.issue}); - return r.setup + iterations(context.stage) * body + r.epilogue; - } -}; - -class SIMTCubeStageCostModel final : public SIMTStageCostModel { -public: - llvm::StringRef getName() const override { return "simt_dot"; } - bool supports(StageCostModelKind kind) const override { - return supportsAny(kind, {StageCostModelKind::CubeRoofline, - StageCostModelKind::TinyCubeRoofline}); - } - double estimate(const StageCostModelContext &context, - const StageImplementation &implementation, - const StageResourceCycles &resources) const override { - const StageResourceCycles &r = resources; - return r.setup + iterations(context.stage) * serialBody(r) + r.epilogue; - } -}; - -class SIMDConversionPackStageCostModel final : public SIMDStageCostModel { -public: - llvm::StringRef getName() const override { return "simd_conversion_pack"; } - bool supports(StageCostModelKind kind) const override { - return kind == StageCostModelKind::ConversionPack; - } - double estimate(const StageCostModelContext &context, - const StageImplementation &implementation, - const StageResourceCycles &resources) const override { - const StageResourceCycles &r = resources; - double body = serialBody(r); - if (permitsSimdOverlap(context.stage)) - body = r.predicate + controlBody(r) + r.spill + - std::max({r.scalar + r.compute, r.load, r.store, r.issue}); - return r.setup + iterations(context.stage) * body + r.epilogue; - } -}; - -class SIMTConversionPackStageCostModel final : public SIMTStageCostModel { -public: - llvm::StringRef getName() const override { return "simt_conversion_pack"; } - bool supports(StageCostModelKind kind) const override { - return kind == StageCostModelKind::ConversionPack; + if (mode == StageMode::SIMD) + return r.setup + count * critical; + const int64_t groups = std::max( + 1, std::min(stage.features.parallelRecurrenceGroupCount, + profile.logicalWarpGroupCount)); + return r.setup + + std::max(std::ceil(count / static_cast(groups)) * critical, + count * r.issue); } - double estimate(const StageCostModelContext &context, - const StageImplementation &implementation, - const StageResourceCycles &resources) const override { - const StageResourceCycles &r = resources; - return r.setup + iterations(context.stage) * serialBody(r) + r.epilogue; + case StageCostModelKind::RowwiseReduction: + return r.setup + + count * std::max(r.scalar + r.load + r.store + r.criticalPath + + controlBody(r) + r.spill, + r.issue); + case StageCostModelKind::CubeRoofline: + case StageCostModelKind::TinyCubeRoofline: + if (mode == StageMode::SIMD && permitsSimdOverlap(stage)) + return r.setup + + count * (r.scalar + r.predicate + controlBody(r) + r.shuffle + + r.spill + + std::max({r.load, r.compute + r.dot, r.store, r.issue})); + return serial; + case StageCostModelKind::ConversionPack: + if (mode == StageMode::SIMD && permitsSimdOverlap(stage)) + return r.setup + count * (r.predicate + controlBody(r) + r.spill + + std::max({r.scalar + r.compute, r.load, r.store, + r.issue})); + return serial; + default: + return serial; } -}; - +} static bool isDeclaredLegal(const LogicalStage &stage, const StageImplementation &implementation) { if (!implementation.isValid()) return false; if (implementation.mode == StageMode::SIMD) - return stage.simdLegal && implementation.superblockFactor == 1; + return stage.simdLegal && implementation.superblockFactor == 1 && + !implementation.localScope; if (!stage.simtLegal) return false; + if (implementation.localScope) + return stage.localSimtMaterializable && + llvm::is_contained(stage.localSimtFactors, + implementation.superblockFactor); return llvm::is_contained(stage.legalSimtFactors, implementation.superblockFactor); } @@ -614,34 +318,6 @@ llvm::StringRef mlir::ascend::stringifyStageCostModel(StageCostModelKind kind) { llvm_unreachable("unknown StageCostModelKind"); } -std::optional -mlir::ascend::parseStageCostModel(llvm::StringRef name) { - return llvm::StringSwitch>(name) - .Case("auto_blockify_dispatch", StageCostModelKind::AutoBlockifyDispatch) - .Case("auto_blockify_loop", StageCostModelKind::AutoBlockifyLoop) - .Case("scalar_issue", StageCostModelKind::ScalarIssue) - .Case("scalar_control", StageCostModelKind::ScalarControl) - .Case("scalar_math", StageCostModelKind::ScalarMath) - .Case("index_generation", StageCostModelKind::IndexGeneration) - .Case("predicate_mask", StageCostModelKind::PredicateMask) - .Case("loop_predicate", StageCostModelKind::LoopPredicate) - .Case("continuous_tile_memory", StageCostModelKind::ContinuousTileMemory) - .Case("continuous_tile_store", StageCostModelKind::ContinuousTileStore) - .Case("continuous_short_load", StageCostModelKind::ContinuousShortLoad) - .Case("cache_policy_store", StageCostModelKind::CachePolicyStore) - .Case("indirect_scalar_memory", StageCostModelKind::IndirectScalarMemory) - .Case("indirect_gather_memory", StageCostModelKind::IndirectGatherMemory) - .Case("independent_pipelined_loop", - StageCostModelKind::IndependentPipelinedLoop) - .Case("loop_carried_recurrence", - StageCostModelKind::LoopCarriedRecurrence) - .Case("rowwise_reduction", StageCostModelKind::RowwiseReduction) - .Case("cube_roofline", StageCostModelKind::CubeRoofline) - .Case("tiny_cube_roofline", StageCostModelKind::TinyCubeRoofline) - .Case("conversion_pack", StageCostModelKind::ConversionPack) - .Default(std::nullopt); -} - bool StageControlFlowRates::isFiniteAndNonNegative() const { const std::array values = { loopBackedgeCycles, conditionalBranchCycles, divergentBranchPenaltyCycles, @@ -697,108 +373,6 @@ bool HardwareProfile::isValid() const { transition.isValid(); } -ProfileProvider::ProfileProvider(HardwareProfile profile) - : profile(std::move(profile)) {} - -llvm::Expected -ProfileProvider::getSnapshot(llvm::StringRef target, - llvm::StringRef profileVersion) const { - if (!profile.isValid()) - return llvm::createStringError(std::errc::invalid_argument, - "hardware profile is invalid"); - if (!target.empty() && target != profile.target) - return llvm::createStringError( - std::errc::invalid_argument, - "hardware profile target '%s' does not match requested target '%s'", - profile.target.c_str(), target.str().c_str()); - if (!profileVersion.empty() && profileVersion != profile.profileVersion) - return llvm::createStringError( - std::errc::invalid_argument, - "hardware profile version '%s' does not match requested version '%s'", - profile.profileVersion.c_str(), profileVersion.str().c_str()); - return &profile; -} - -StageCostModelRegistry::StageCostModelRegistry() { - models.push_back(std::make_unique()); - models.push_back(std::make_unique()); - models.push_back(std::make_unique()); - models.push_back(std::make_unique()); - models.push_back(std::make_unique()); - models.push_back(std::make_unique()); - models.push_back(std::make_unique()); - models.push_back(std::make_unique()); - models.push_back(std::make_unique()); - models.push_back(std::make_unique()); - models.push_back(std::make_unique()); - models.push_back(std::make_unique()); - models.push_back(std::make_unique()); - models.push_back(std::make_unique()); - models.push_back(std::make_unique()); - models.push_back(std::make_unique()); - models.push_back(std::make_unique()); - models.push_back(std::make_unique()); -} - -const StageCostModelRegistry &StageCostModelRegistry::get() { - static const StageCostModelRegistry registry; - return registry; -} - -llvm::Expected -StageCostModelRegistry::lookup(StageMode mode, StageCostModelKind kind) const { - const StageCostModel *match = nullptr; - for (const std::unique_ptr &model : models) { - if (model->getMode() != mode || !model->supports(kind)) - continue; - if (match) - return llvm::createStringError( - std::errc::invalid_argument, - "multiple StageCostModels registered for (%s, %s)", - stringifyStageMode(mode).str().c_str(), - stringifyStageCostModel(kind).str().c_str()); - match = model.get(); - } - if (!match) - return llvm::createStringError(std::errc::invalid_argument, - "no StageCostModel registered for (%s, %s)", - stringifyStageMode(mode).str().c_str(), - stringifyStageCostModel(kind).str().c_str()); - return match; -} - -llvm::Error StageCostModelRegistry::verifyComplete() const { - constexpr std::array kinds = { - StageCostModelKind::AutoBlockifyDispatch, - StageCostModelKind::AutoBlockifyLoop, - StageCostModelKind::ScalarIssue, - StageCostModelKind::ScalarControl, - StageCostModelKind::ScalarMath, - StageCostModelKind::IndexGeneration, - StageCostModelKind::PredicateMask, - StageCostModelKind::LoopPredicate, - StageCostModelKind::ContinuousTileMemory, - StageCostModelKind::ContinuousTileStore, - StageCostModelKind::ContinuousShortLoad, - StageCostModelKind::CachePolicyStore, - StageCostModelKind::IndirectScalarMemory, - StageCostModelKind::IndirectGatherMemory, - StageCostModelKind::IndependentPipelinedLoop, - StageCostModelKind::LoopCarriedRecurrence, - StageCostModelKind::RowwiseReduction, - StageCostModelKind::CubeRoofline, - StageCostModelKind::TinyCubeRoofline, - StageCostModelKind::ConversionPack}; - for (StageMode mode : {StageMode::SIMD, StageMode::SIMT}) { - for (StageCostModelKind kind : kinds) { - auto model = lookup(mode, kind); - if (!model) - return model.takeError(); - } - } - return llvm::Error::success(); -} - llvm::Expected StageCostEvaluator::evaluate(const StagePartition &partition, const HardwareProfile &profile) const { @@ -809,12 +383,8 @@ StageCostEvaluator::evaluate(const StagePartition &partition, if (!profile.isValid()) return llvm::createStringError(std::errc::invalid_argument, "HardwareProfile is invalid"); - if (llvm::Error error = registry.verifyComplete()) - return std::move(error); - StageCostTable table; table.domain = partition.domain; - table.boundarySource = partition.boundarySource; table.operationOwnershipComplete = partition.operationOwnershipComplete; table.modeledOperationCount = partition.modeledOperationCount; table.profileVersion = profile.profileVersion; @@ -826,7 +396,6 @@ StageCostEvaluator::evaluate(const StagePartition &partition, "every Phase requires an id and Stage"); LogicalPhaseCost phaseCost; phaseCost.id = phase.id; - phaseCost.description = phase.description; for (const LogicalStage &stage : phase.stages) { if (stage.id.empty() || !stageIds.insert(stage.id).second) @@ -849,7 +418,6 @@ StageCostEvaluator::evaluate(const StagePartition &partition, LogicalStageCost logicalCost; logicalCost.id = stage.id; - logicalCost.description = stage.description; logicalCost.model = stringifyStageCostModel(stage.costModelKind).str(); logicalCost.schedule = stage.scheduleKind; logicalCost.iterationCount = stage.iterationCount; @@ -866,39 +434,35 @@ StageCostEvaluator::evaluate(const StagePartition &partition, logicalCost.scopeOutputTensorBytes = stage.scopeOutputTensorBytes; logicalCost.simtAnchorIndices = stage.simtAnchorIndices; logicalCost.localSimtMaterializable = stage.localSimtMaterializable; + logicalCost.legalSimtFactors = stage.legalSimtFactors; logicalCost.localSimtFactors = stage.localSimtFactors; llvm::SmallVector implementations; if (stage.simdLegal) - implementations.push_back({StageMode::SIMD, 1}); + implementations.push_back({StageMode::SIMD, 1, false}); if (stage.simtLegal) for (int64_t factor : stage.legalSimtFactors) - implementations.push_back({StageMode::SIMT, factor}); + implementations.push_back({StageMode::SIMT, factor, false}); + if (stage.simtLegal && stage.localSimtMaterializable) + for (int64_t factor : stage.localSimtFactors) + implementations.push_back({StageMode::SIMT, factor, true}); for (const StageImplementation &implementation : implementations) { if (!isDeclaredLegal(stage, implementation)) return llvm::createStringError(std::errc::invalid_argument, "Stage '%s' has an illegal candidate", stage.id.c_str()); - auto model = registry.lookup(implementation.mode, stage.costModelKind); - if (!model) - return model.takeError(); StageResourceCycles resources = - implementation.mode == StageMode::SIMD - ? mapSIMDWorkload(stage, profile.simd) - : mapSIMTWorkload(stage, profile.simt); - const StageCostModelContext context{stage, profile}; - + mapWorkload(stage, + implementation.mode == StageMode::SIMD ? profile.simd + : profile.simt, + implementation.mode); StageImplementationCost cost; cost.implementation = implementation; cost.resources = resources; - cost.modelName = (*model)->getName().str(); - cost.profileVersion = profile.profileVersion; - cost.source = - "post-transform TTIR StageWorkload + immutable HardwareProfile"; cost.totalCycles = applySuperBlock( stage, resources, implementation, profile, - (*model)->estimate(context, implementation, resources)); + estimateStage(stage, profile, implementation.mode, resources)); if (!cost.isValid()) return llvm::createStringError(std::errc::invalid_argument, "Stage '%s' produced an invalid cost", diff --git a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageRouteCostModel.cpp b/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageRouteCostModel.cpp index 052d7e97d0..ac65f900a1 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageRouteCostModel.cpp +++ b/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageRouteCostModel.cpp @@ -8,8 +8,7 @@ #include #include #include -#include -#include +#include #include using namespace mlir; @@ -17,38 +16,25 @@ using namespace mlir::ascend; namespace { -enum class RouteClass : unsigned { AllSIMD = 0, AllSIMT = 1, Mixed = 2 }; - -struct PartialRoute { - double totalCycles = 0.0; - /// Counterfactual cost used if this prefix later becomes a mixed route. - /// Local SIMT Stages retain their selected F1/F2/F4 implementation and add - /// the exact UB-backed scope-boundary cost. - double mixedEquivalentCycles = 0.0; - StageMode exitMode = StageMode::SIMD; - RouteClass routeClass = RouteClass::AllSIMD; - bool allSimtStagesLocal = true; - std::vector implementations; - std::vector entryTransitionCycles; - std::vector logicalStageCycles; - std::vector mixedEquivalentStageCycles; - int64_t routeSuperblockFactor = 1; -}; - static double mixedEquivalentStageCost(const LogicalStageCost &stage, const StageImplementationCost &selected, const StageTransitionCost &transition) { if (selected.implementation.mode != StageMode::SIMT || - !stage.localSimtMaterializable) + !selected.implementation.localScope || !stage.localSimtMaterializable) return selected.totalCycles; // Materializer currently creates one scope per primitive anchor. The // route DP otherwise observes only one Stage-mode change and would charge // one transition pair even when the generated TTIR contains several local - // scopes. Charge the additional physical pairs here; all-SIMT routes do - // not consume this mixed-only equivalent cost. + // scopes. Scope-local SuperBlock groups F independent logical programs in + // one outlined SIMT VF, so those programs share each fixed SIMD/SIMT mode + // switch. The fixed transition cost is therefore amortized by F. Tensor + // handoff bytes are not divided: the materializer still transfers every + // logical program's live-in/live-out values through UB. const int64_t scopeCount = std::max(1, stage.localSimtScopeCount); + const double factor = static_cast( + std::max(1, selected.implementation.superblockFactor)); const double fixedScopeTransitions = - static_cast(scopeCount) * + static_cast(scopeCount) / factor * (transition.get(StageMode::SIMD, StageMode::SIMT) + transition.get(StageMode::SIMT, StageMode::SIMD)); const double activeThreads = @@ -72,37 +58,24 @@ static double mixedEquivalentStageCost(const LogicalStageCost &stage, outputHandoffCycles; } -static unsigned modeIndex(StageMode mode) { - return mode == StageMode::SIMD ? 0u : 1u; -} - -static RouteClass initialClass(StageMode mode) { - return mode == StageMode::SIMD ? RouteClass::AllSIMD : RouteClass::AllSIMT; -} - -static RouteClass appendClass(RouteClass current, StageMode next) { - if (current == RouteClass::Mixed) - return current; - if ((current == RouteClass::AllSIMD && next == StageMode::SIMD) || - (current == RouteClass::AllSIMT && next == StageMode::SIMT)) - return current; - return RouteClass::Mixed; -} - -static StageRoutePlan toPlan(const std::optional &route, - StageKernelRouteKind kind) { - StageRoutePlan result; - result.candidate = kind; - if (!route) - return result; - result.legal = true; - result.implementations = route->implementations; - result.entryTransitionCycles = route->entryTransitionCycles; - result.logicalStageCycles = route->logicalStageCycles; - result.routeSuperblockFactor = route->routeSuperblockFactor; - result.totalCycles = route->totalCycles; - result.source = "stage_dynamic_programming"; - return result; +/// AutoBlockify V1 is a route-conditional execution schedule. The analysis +/// view contains its real dispatch/loop operations so pure-SIMT and Mixed can +/// pay them, but an all-SIMD executable restores the original logical grid. +/// Keep the Stage positions for report alignment and remove only their cost +/// from the all-SIMD candidate. +static void removeAutoBlockifyCostFromAllSIMD(StageRoutePlan &plan, + const StageCostTable &costTable) { + if (!plan.legal || plan.logicalStageCycles.size() != costTable.stages.size()) + return; + for (size_t index = 0; index < costTable.stages.size(); ++index) { + const llvm::StringRef model = costTable.stages[index].model; + if (model != "auto_blockify_dispatch" && model != "auto_blockify_loop") + continue; + plan.totalCycles -= plan.logicalStageCycles[index]; + plan.logicalStageCycles[index] = 0.0; + plan.entryTransitionCycles[index] = 0.0; + } + plan.totalCycles = std::max(0.0, plan.totalCycles); } } // namespace @@ -111,8 +84,7 @@ llvm::StringRef mlir::ascend::stringifyStageMode(StageMode mode) { return mode == StageMode::SIMD ? "simd" : "simt"; } -llvm::StringRef -mlir::ascend::stringifyStageKernelRoute(StageKernelRouteKind kind) { +static llvm::StringRef stringifyStageKernelRoute(StageKernelRouteKind kind) { switch (kind) { case StageKernelRouteKind::AllSIMD: return "all_simd"; @@ -124,7 +96,7 @@ mlir::ascend::stringifyStageKernelRoute(StageKernelRouteKind kind) { llvm_unreachable("unknown stage kernel route kind"); } -llvm::StringRef mlir::ascend::stringifyStageSchedule(StageScheduleKind kind) { +static llvm::StringRef stringifyStageSchedule(StageScheduleKind kind) { switch (kind) { case StageScheduleKind::StraightLine: return "straight_line"; @@ -141,14 +113,17 @@ llvm::StringRef mlir::ascend::stringifyStageSchedule(StageScheduleKind kind) { bool StageImplementation::isValid() const { if (superblockFactor <= 0 || (superblockFactor & (superblockFactor - 1)) != 0) return false; - return mode == StageMode::SIMT || superblockFactor == 1; + if (mode == StageMode::SIMD) + return superblockFactor == 1 && !localScope; + return true; } llvm::json::Object StageImplementation::toJSON() const { - llvm::json::Object result; - result["mode"] = stringifyStageMode(mode); - result["superblock_factor"] = superblockFactor; - return result; + return llvm::json::Object{ + {"mode", stringifyStageMode(mode)}, + {"superblock_factor", superblockFactor}, + {"materialization", + localScope ? "local_simt_scope_with_kernel_v1" : "whole_kernel"}}; } bool StageModelFeatures::isValid() const { @@ -221,15 +196,14 @@ llvm::json::Object StageModelFeatures::toJSON() const { result["parallel_recurrence_group_count"] = parallelRecurrenceGroupCount; result["active_lane_ratio"] = activeLaneRatio; result["simd_roofline_permitted"] = permitsSimdRoofline(); - result["source"] = source; return result; } bool StageResourceCycles::isFiniteAndNonNegative() const { - const std::array values = { - setup, scalar, load, store, compute, predicate, - shuffle, dot, control, loopControl, branchControl, divergence, - synchronization, spill, issue, criticalPath, epilogue}; + const std::array values = { + setup, scalar, load, store, compute, + predicate, shuffle, dot, loopControl, branchControl, + divergence, synchronization, spill, issue, criticalPath}; return std::all_of(values.begin(), values.end(), [](double value) { return std::isfinite(value) && value >= 0.0; }); @@ -245,7 +219,6 @@ llvm::json::Object StageResourceCycles::toJSON() const { result["predicate_per_iteration"] = predicate; result["shuffle_per_iteration"] = shuffle; result["dot_per_iteration"] = dot; - result["control_per_iteration"] = control; result["loop_control_per_iteration"] = loopControl; result["branch_control_per_iteration"] = branchControl; result["divergence_per_iteration"] = divergence; @@ -253,31 +226,23 @@ llvm::json::Object StageResourceCycles::toJSON() const { result["spill_per_iteration"] = spill; result["issue_per_iteration"] = issue; result["critical_path_per_iteration"] = criticalPath; - result["epilogue"] = epilogue; return result; } bool StageImplementationCost::isValid() const { return implementation.isValid() && std::isfinite(totalCycles) && - totalCycles >= 0.0 && resources.isFiniteAndNonNegative() && - !modelName.empty() && !profileVersion.empty(); + totalCycles >= 0.0 && resources.isFiniteAndNonNegative(); } llvm::json::Object StageImplementationCost::toJSON() const { - llvm::json::Object result; - result["implementation"] = implementation.toJSON(); - result["total_system_cycles"] = totalCycles; - result["resource_system_cycles"] = resources.toJSON(); - result["model_name"] = modelName; - result["profile_version"] = profileVersion; - result["source"] = source; - return result; + return llvm::json::Object{{"implementation", implementation.toJSON()}, + {"total_system_cycles", totalCycles}, + {"resource_system_cycles", resources.toJSON()}}; } llvm::json::Object LogicalStageCost::toJSON() const { llvm::json::Object result; result["id"] = id; - result["description"] = description; result["model"] = model; result["schedule_kind"] = stringifyStageSchedule(schedule); result["iteration_count"] = iterationCount; @@ -296,6 +261,10 @@ llvm::json::Object LogicalStageCost::toJSON() const { anchorIndices.push_back(static_cast(index)); result["simt_anchor_indices"] = std::move(anchorIndices); result["local_simt_materializable"] = localSimtMaterializable; + llvm::json::Array legalFactors; + for (int64_t factor : legalSimtFactors) + legalFactors.push_back(factor); + result["legal_simt_factors"] = std::move(legalFactors); llvm::json::Array localFactors; for (int64_t factor : localSimtFactors) localFactors.push_back(factor); @@ -308,28 +277,10 @@ llvm::json::Object LogicalStageCost::toJSON() const { } llvm::json::Object LogicalPhaseCost::toJSON() const { - llvm::json::Object result; - result["id"] = id; - result["description"] = description; llvm::json::Array stageArray; for (const LogicalStageCost &stage : stages) stageArray.push_back(stage.toJSON()); - result["stages"] = std::move(stageArray); - return result; -} - -llvm::json::Object StageCostTable::toJSON() const { - llvm::json::Object result; - result["domain"] = domain; - result["boundary_source"] = boundarySource; - result["operation_ownership_complete"] = operationOwnershipComplete; - result["modeled_operation_count"] = modeledOperationCount; - result["profile_version"] = profileVersion; - llvm::json::Array phaseArray; - for (const LogicalPhaseCost &phase : phases) - phaseArray.push_back(phase.toJSON()); - result["phases"] = std::move(phaseArray); - return result; + return llvm::json::Object{{"id", id}, {"stages", std::move(stageArray)}}; } bool StageTransitionCost::isValid() const { @@ -362,7 +313,6 @@ llvm::json::Object StageTransitionCost::toJSON() const { result["simt_ub_store_bytes_per_thread_per_system_cycle"] = simtUbStoreBytesPerThreadPerCycle; result["simt_warp_size"] = simtWarpSize; - result["source"] = source; return result; } @@ -372,7 +322,8 @@ llvm::json::Object StageRoutePlan::toJSON() const { result["legal"] = legal; result["total_system_cycles"] = totalCycles; result["route_superblock_factor"] = routeSuperblockFactor; - result["source"] = source; + result["runtime_physical_program_count"] = runtimePhysicalProgramCount; + result["runtime_wave_count"] = runtimeWaveCount; llvm::json::Array stages; for (size_t i = 0; i < implementations.size(); ++i) { llvm::json::Object stage; @@ -393,7 +344,7 @@ llvm::json::Object StageCostModelSummary::toJSON() const { llvm::json::Object result; result["applied"] = applied; result["domain"] = domain; - result["boundary_source"] = boundarySource; + result["boundary_source"] = "operation_graph"; result["operation_ownership_complete"] = operationOwnershipComplete; result["modeled_operation_count"] = modeledOperationCount; result["profile_version"] = profileVersion; @@ -426,126 +377,142 @@ mlir::ascend::solveStageRoutes(const StageCostTable &costTable, "stage transition costs must be finite and " "non-negative"); - // Keep one best partial route for every (exit mode, route class, - // whole-kernel SIMT SuperBlock factor). Collapsing the factor dimension - // can discard a slightly slower F1 prefix that becomes globally optimal, - // or worse, combine F1 and F4 Stage costs into an unrealizable F4 kernel. - using FactorRoutes = std::map; - using State = std::array, 2>; - State current; - bool firstStage = true; - for (const LogicalStageCost &stage : costTable.stages) { - State next; - if (stage.implementations.empty()) - return llvm::createStringError(std::errc::invalid_argument, - "stage '%s' has no legal implementation", - stage.id.c_str()); - - for (const StageImplementationCost &cost : stage.implementations) { - if (!cost.isValid()) - return llvm::createStringError( - std::errc::invalid_argument, - "stage '%s' has an invalid implementation cost", stage.id.c_str()); - - if (firstStage) { - PartialRoute route; - route.totalCycles = cost.totalCycles; - route.mixedEquivalentCycles = - mixedEquivalentStageCost(stage, cost, transition); - route.exitMode = cost.implementation.mode; - route.routeClass = initialClass(cost.implementation.mode); - route.allSimtStagesLocal = - cost.implementation.mode != StageMode::SIMT || - (stage.localSimtMaterializable && - llvm::is_contained(stage.localSimtFactors, - cost.implementation.superblockFactor)); - route.implementations.push_back(cost.implementation); - route.routeSuperblockFactor = cost.implementation.superblockFactor; - route.entryTransitionCycles.push_back(0.0); - route.logicalStageCycles.push_back(cost.totalCycles); - route.mixedEquivalentStageCycles.push_back(route.mixedEquivalentCycles); - auto &routes = next[modeIndex(route.exitMode)] - [static_cast(route.routeClass)]; - auto [slot, inserted] = - routes.try_emplace(route.routeSuperblockFactor, route); - if (!inserted && route.totalCycles < slot->second.totalCycles) - slot->second = std::move(route); - continue; - } + auto findImplementation = + [](const LogicalStageCost &stage, StageMode mode, int64_t factor, + bool localScope) -> const StageImplementationCost * { + for (const StageImplementationCost &cost : stage.implementations) + if (cost.implementation.mode == mode && + cost.implementation.superblockFactor == factor && + cost.implementation.localScope == localScope) + return &cost; + return nullptr; + }; - for (const auto &byClass : current) { - for (const auto &factorRoutes : byClass) { - for (const auto &factorRoute : factorRoutes) { - const PartialRoute &previous = factorRoute.second; - const bool routeAlreadyHasSimt = - previous.routeClass != RouteClass::AllSIMD; - if (cost.implementation.mode == StageMode::SIMT && - routeAlreadyHasSimt && - cost.implementation.superblockFactor != - previous.routeSuperblockFactor) - continue; - - // Local scopes pay both physical directions in - // mixedEquivalentStageCost. Adjacent Stage labels are a logical - // route description, not an additional hardware transition. - const double entryTransition = 0.0; - PartialRoute route = previous; - RouteClass nextClass = - appendClass(route.routeClass, cost.implementation.mode); - route.routeClass = nextClass; - route.allSimtStagesLocal = - route.allSimtStagesLocal && - (cost.implementation.mode != StageMode::SIMT || - (stage.localSimtMaterializable && - llvm::is_contained(stage.localSimtFactors, - cost.implementation.superblockFactor))); - if (route.routeClass == RouteClass::Mixed && - !route.allSimtStagesLocal) - continue; - route.exitMode = cost.implementation.mode; - route.implementations.push_back(cost.implementation); - if (cost.implementation.mode == StageMode::SIMT && - !routeAlreadyHasSimt) - route.routeSuperblockFactor = - cost.implementation.superblockFactor; - route.entryTransitionCycles.push_back(entryTransition); - const double logicalStageCycles = - entryTransition + cost.totalCycles; - const double mixedLogicalStageCycles = - entryTransition + - mixedEquivalentStageCost(stage, cost, transition); - route.mixedEquivalentCycles += mixedLogicalStageCycles; - route.mixedEquivalentStageCycles.push_back(mixedLogicalStageCycles); - if (nextClass == RouteClass::Mixed) { - route.totalCycles = route.mixedEquivalentCycles; - route.logicalStageCycles = route.mixedEquivalentStageCycles; - } else { - route.logicalStageCycles.push_back(logicalStageCycles); - route.totalCycles += logicalStageCycles; - } - auto &routes = next[modeIndex(route.exitMode)] - [static_cast(route.routeClass)]; - auto [slot, inserted] = - routes.try_emplace(route.routeSuperblockFactor, route); - if (!inserted && route.totalCycles < slot->second.totalCycles) - slot->second = std::move(route); + auto buildPlan = [&](StageKernelRouteKind kind, + int64_t factor) -> StageRoutePlan { + StageRoutePlan plan; + plan.candidate = kind; + plan.routeSuperblockFactor = factor; + struct MixedChoice { + const StageImplementationCost *simd = nullptr; + const StageImplementationCost *simt = nullptr; + double simdCycles = std::numeric_limits::infinity(); + double simtCycles = std::numeric_limits::infinity(); + }; + std::vector mixedChoices; + mixedChoices.reserve(costTable.stages.size()); + for (const LogicalStageCost &stage : costTable.stages) { + for (const StageImplementationCost &cost : stage.implementations) + if (!cost.isValid()) + return plan; + + const StageImplementationCost *selected = nullptr; + double stageCycles = 0.0; + if (kind == StageKernelRouteKind::AllSIMD) { + selected = findImplementation(stage, StageMode::SIMD, 1, false); + } else if (kind == StageKernelRouteKind::AllSIMT) { + selected = findImplementation(stage, StageMode::SIMT, factor, false); + } else { + const StageImplementationCost *simd = + findImplementation(stage, StageMode::SIMD, 1, false); + const StageImplementationCost *simt = + findImplementation(stage, StageMode::SIMT, factor, true); + const double simdCycles = + simd ? simd->totalCycles : std::numeric_limits::infinity(); + const double simtCycles = + simt ? mixedEquivalentStageCost(stage, *simt, transition) + : std::numeric_limits::infinity(); + selected = simtCycles < simdCycles ? simt : simd; + stageCycles = std::min(simdCycles, simtCycles); + mixedChoices.push_back({simd, simt, simdCycles, simtCycles}); + } + if (!selected) + return plan; + if (kind != StageKernelRouteKind::Mixed) + stageCycles = selected->totalCycles; + + const double transitionCycles = stageCycles - selected->totalCycles; + plan.implementations.push_back(selected->implementation); + plan.entryTransitionCycles.push_back(transitionCycles); + plan.logicalStageCycles.push_back(stageCycles); + plan.totalCycles += stageCycles; + } + if (kind == StageKernelRouteKind::Mixed) { + auto countMode = [&](StageMode mode) { + return llvm::count_if(plan.implementations, [&](const auto &selected) { + return selected.mode == mode; + }); + }; + auto forceOneMode = [&](StageMode required) { + size_t bestIndex = plan.implementations.size(); + double bestPenalty = std::numeric_limits::infinity(); + for (size_t index = 0; index < mixedChoices.size(); ++index) { + const MixedChoice &choice = mixedChoices[index]; + const StageImplementationCost *replacement = + required == StageMode::SIMD ? choice.simd : choice.simt; + if (!replacement) + continue; + const double replacementCycles = required == StageMode::SIMD + ? choice.simdCycles + : choice.simtCycles; + const double penalty = + replacementCycles - plan.logicalStageCycles[index]; + if (penalty < bestPenalty) { + bestPenalty = penalty; + bestIndex = index; } } + if (bestIndex == plan.implementations.size()) + return false; + const MixedChoice &choice = mixedChoices[bestIndex]; + const StageImplementationCost *replacement = + required == StageMode::SIMD ? choice.simd : choice.simt; + const double replacementCycles = + required == StageMode::SIMD ? choice.simdCycles : choice.simtCycles; + plan.totalCycles += + replacementCycles - plan.logicalStageCycles[bestIndex]; + plan.implementations[bestIndex] = replacement->implementation; + plan.logicalStageCycles[bestIndex] = replacementCycles; + plan.entryTransitionCycles[bestIndex] = + replacementCycles - replacement->totalCycles; + return true; + }; + + // Mixed is a constrained candidate, not the unconstrained per-Stage + // minimum. If the latter collapses to all-SIMD or all-SIMT, switch the + // Stage with the smallest incremental cost so the reported candidate + // is the cheapest route that genuinely contains both modes. + if ((countMode(StageMode::SIMT) == 0 && !forceOneMode(StageMode::SIMT)) || + (countMode(StageMode::SIMD) == 0 && !forceOneMode(StageMode::SIMD)) || + countMode(StageMode::SIMD) == 0 || countMode(StageMode::SIMT) == 0) { + StageRoutePlan invalid; + invalid.candidate = kind; + return invalid; } } - current = std::move(next); - firstStage = false; - } + if (costTable.logicalProgramCountHint > 0) { + plan.runtimePhysicalProgramCount = + (costTable.logicalProgramCountHint + factor - 1) / factor; + if (costTable.physicalCoreCountHint > 0) + plan.runtimeWaveCount = (plan.runtimePhysicalProgramCount + + costTable.physicalCoreCountHint - 1) / + costTable.physicalCoreCountHint; + for (double &cycles : plan.logicalStageCycles) + cycles *= static_cast(plan.runtimeWaveCount); + plan.totalCycles *= static_cast(plan.runtimeWaveCount); + } + plan.legal = true; + return plan; + }; - auto bestClass = [&](RouteClass routeClass) -> std::optional { - std::optional best; - for (const auto &byClass : current) { - const auto &candidates = byClass[static_cast(routeClass)]; - for (const auto &factorRoute : candidates) { - const PartialRoute &candidate = factorRoute.second; - if (!best || candidate.totalCycles < best->totalCycles) - best = candidate; - } + auto bestFactoredPlan = [&](StageKernelRouteKind kind) { + StageRoutePlan best; + best.candidate = kind; + for (int64_t factor : {1, 2, 4}) { + StageRoutePlan candidate = buildPlan(kind, factor); + if (candidate.legal && + (!best.legal || candidate.totalCycles < best.totalCycles)) + best = std::move(candidate); } return best; }; @@ -553,19 +520,16 @@ mlir::ascend::solveStageRoutes(const StageCostTable &costTable, StageCostModelSummary result; result.applied = true; result.domain = costTable.domain; - result.boundarySource = costTable.boundarySource; result.operationOwnershipComplete = costTable.operationOwnershipComplete; result.modeledOperationCount = costTable.modeledOperationCount; result.profileVersion = costTable.profileVersion; result.phases = costTable.phases; result.stages = costTable.stages; result.transition = transition; - result.allSimd = - toPlan(bestClass(RouteClass::AllSIMD), StageKernelRouteKind::AllSIMD); - result.allSimt = - toPlan(bestClass(RouteClass::AllSIMT), StageKernelRouteKind::AllSIMT); - result.mixed = - toPlan(bestClass(RouteClass::Mixed), StageKernelRouteKind::Mixed); + result.allSimd = buildPlan(StageKernelRouteKind::AllSIMD, 1); + result.allSimt = bestFactoredPlan(StageKernelRouteKind::AllSIMT); + result.mixed = bestFactoredPlan(StageKernelRouteKind::Mixed); + removeAutoBlockifyCostFromAllSIMD(result.allSimd, costTable); auto aggregatePhases = [&](StageRoutePlan &plan) { if (!plan.legal) diff --git a/third_party/ascend/costmodel/lib/AscendModel/Transforms/CMakeLists.txt b/third_party/ascend/costmodel/lib/AscendModel/Transforms/CMakeLists.txt index fb2288e354..2295c1ddfc 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/Transforms/CMakeLists.txt +++ b/third_party/ascend/costmodel/lib/AscendModel/Transforms/CMakeLists.txt @@ -1,4 +1,5 @@ add_mlir_library(AscendModelTransforms + PARTIAL_SOURCES_INTENDED ConvertTritonToAscend.cpp InsertDataTransfers.cpp AssignOpIDs.cpp diff --git a/third_party/ascend/costmodel/lib/AscendModel/Transforms/HIVMAnalysisPass.cpp b/third_party/ascend/costmodel/lib/AscendModel/Transforms/HIVMAnalysisPass.cpp index 7e8ab2f26f..e55fd0b19e 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/Transforms/HIVMAnalysisPass.cpp +++ b/third_party/ascend/costmodel/lib/AscendModel/Transforms/HIVMAnalysisPass.cpp @@ -22,14 +22,6 @@ namespace ascend { namespace { -static FailureOr parseSchedulerMode(llvm::StringRef mode) { - if (mode.empty() || mode == "static") - return HIVMSchedulerMode::Static; - if (mode == "des") - return HIVMSchedulerMode::DES; - return failure(); -} - struct HIVMAnalysisPass : public impl::HIVMAnalysisPassBase { using HIVMAnalysisPassBase::HIVMAnalysisPassBase; @@ -45,14 +37,7 @@ struct HIVMAnalysisPass : public impl::HIVMAnalysisPassBase { return; } const HardwareConfig &config = *hardwareConfig; - auto schedulerOr = parseSchedulerMode(schedulerMode); - if (failed(schedulerOr)) { - module.emitError() << "invalid HIVM scheduler mode `" << schedulerMode - << "`; expected `static` or `des`"; - signalPassFailure(); - return; - } - HIVMAnalyzer analyzer(config, argBindingsStr, *schedulerOr); + HIVMAnalyzer analyzer(config, argBindingsStr); HIVMAnalysisReport report; std::string error; if (!analyzer.analyzeModule(module, report, error)) { diff --git a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/Transforms/MaterializeSimtScopes.cpp b/third_party/ascend/costmodel/lib/AscendModel/Transforms/MaterializeSimtScopes.cpp similarity index 85% rename from third_party/ascend/costmodel/lib/AscendModel/RouteModel/Transforms/MaterializeSimtScopes.cpp rename to third_party/ascend/costmodel/lib/AscendModel/Transforms/MaterializeSimtScopes.cpp index e701061cb3..3ec66db58c 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/Transforms/MaterializeSimtScopes.cpp +++ b/third_party/ascend/costmodel/lib/AscendModel/Transforms/MaterializeSimtScopes.cpp @@ -7,9 +7,9 @@ // //===----------------------------------------------------------------------===// -#include "AscendModel/RouteModel/SimtAnchorAnalysis.h" -#include "AscendModel/RouteModel/SimtSelection.h" +#include "AscendModel/Analysis/SimtAnchorAnalysis.h" #include "AscendModel/Transforms/Passes.h" +#include "AscendModel/Transforms/SimtSelection.h" #include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinAttributes.h" @@ -48,13 +48,24 @@ static bool isMaterializable(Operation *op) { /// Scope regions are not isolated from above, so operands remain legal /// captures. Moving only the planned operation keeps SIMD producers and /// consumers outside the SIMT region. -static LogicalResult wrapAnchorOperation(Operation *op) { +inline constexpr llvm::StringLiteral kScopeSuperblockFactorAttr = + "ascend.scope_superblock.factor"; + +static void setScopeExecutionAttrs(Operation *scopeOp, OpBuilder &builder, + int64_t superblockFactor) { + scopeOp->setAttr(kVectorModeAttr, builder.getStringAttr("simt")); + scopeOp->setAttr(kScopeSuperblockFactorAttr, + builder.getI64IntegerAttr(superblockFactor)); +} + +static LogicalResult wrapAnchorOperation(Operation *op, + int64_t superblockFactor) { OpBuilder builder(op); OperationState scopeState(op->getLoc(), "scope.scope"); scopeState.addTypes(op->getResultTypes()); - scopeState.addAttribute(kVectorModeAttr, builder.getStringAttr("simt")); scopeState.addRegion(); Operation *scopeOp = builder.create(scopeState); + setScopeExecutionAttrs(scopeOp, builder, superblockFactor); Region &scopeRegion = scopeOp->getRegion(0); auto *scopeBody = new Block(); @@ -82,7 +93,8 @@ static LogicalResult wrapAnchorOperation(Operation *op) { /// `insertionPoint` lets solve_tril move pure mask setup across the initial /// loads while keeping those loads outside, matching the hand-written scope. static LogicalResult wrapAnchorRange(ArrayRef ops, - Operation *insertionPoint) { + Operation *insertionPoint, + int64_t superblockFactor) { if (ops.empty()) return success(); Block *parent = insertionPoint ? insertionPoint->getBlock() : nullptr; @@ -120,9 +132,9 @@ static LogicalResult wrapAnchorRange(ArrayRef ops, for (Value value : escaping) escapingTypes.push_back(value.getType()); scopeState.addTypes(escapingTypes); - scopeState.addAttribute(kVectorModeAttr, builder.getStringAttr("simt")); scopeState.addRegion(); Operation *scopeOp = builder.create(scopeState); + setScopeExecutionAttrs(scopeOp, builder, superblockFactor); Region &scopeRegion = scopeOp->getRegion(0); auto *scopeBody = new Block(); @@ -149,7 +161,11 @@ static LogicalResult wrapAnchorRange(ArrayRef ops, } // namespace LogicalResult materializeSimtAnchorPlan(ModuleOp module, - const SimtAnchorPlan &plan) { + const SimtAnchorPlan &plan, + int64_t superblockFactor) { + if (superblockFactor <= 0 || (superblockFactor & (superblockFactor - 1)) != 0) + return module.emitError( + "SIMT scope superblock factor must be a positive power of two"); struct PlannedRange { SmallVector operations; Operation *insertionPoint = nullptr; @@ -183,12 +199,13 @@ LogicalResult materializeSimtAnchorPlan(ModuleOp module, int64_t materialized = 0; for (const PlannedRange &range : anchorRanges) { - if (failed(wrapAnchorRange(range.operations, range.insertionPoint))) + if (failed(wrapAnchorRange(range.operations, range.insertionPoint, + superblockFactor))) return failure(); ++materialized; } for (Operation *op : anchorOps) { - if (failed(wrapAnchorOperation(op))) + if (failed(wrapAnchorOperation(op, superblockFactor))) return failure(); ++materialized; } diff --git a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/Transforms/SelectSimdSimtCostModel.cpp b/third_party/ascend/costmodel/lib/AscendModel/Transforms/SelectSimdSimtCostModel.cpp similarity index 62% rename from third_party/ascend/costmodel/lib/AscendModel/RouteModel/Transforms/SelectSimdSimtCostModel.cpp rename to third_party/ascend/costmodel/lib/AscendModel/Transforms/SelectSimdSimtCostModel.cpp index 301772a4f3..cdf66e6b27 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/Transforms/SelectSimdSimtCostModel.cpp +++ b/third_party/ascend/costmodel/lib/AscendModel/Transforms/SelectSimdSimtCostModel.cpp @@ -7,15 +7,16 @@ // //===----------------------------------------------------------------------===// +#include "AscendModel/Analysis/SimtAnchorAnalysis.h" #include "AscendModel/RouteModel/SimdSimtCostModel.h" -#include "AscendModel/RouteModel/SimtAnchorAnalysis.h" -#include "AscendModel/RouteModel/SimtSelection.h" #include "AscendModel/Transforms/Passes.h" +#include "AscendModel/Transforms/SimtSelection.h" #include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/Operation.h" +#include "mlir/Parser/Parser.h" #include "mlir/Pass/Pass.h" #include "llvm/ADT/DenseSet.h" #include "llvm/ADT/SmallVector.h" @@ -95,16 +96,32 @@ buildSelectedMixedAnchorPlan(const StageCostModelSummary &stageModel, if (implementation.mode != StageMode::SIMT) continue; - for (unsigned index : stage.simtAnchorIndices) { - if (index >= completePlan.anchors.size() || - !included.insert(index).second) - continue; - selected.anchors.push_back(completePlan.anchors[index]); - } + llvm::SmallVector stageAnchorIndices; + for (unsigned index : stage.simtAnchorIndices) + if (index < completePlan.anchors.size() && included.insert(index).second) + stageAnchorIndices.push_back(index); + auto merged = mergeSimtStageAnchors(completePlan, stageAnchorIndices); + if (!stageAnchorIndices.empty() && !merged) + return SimtAnchorPlan{}; + if (merged) + selected.anchors.push_back(std::move(*merged)); } return selected; } +static bool +anchorPlansHaveCompatibleIndices(const SimtAnchorPlan &analysis, + const SimtAnchorPlan &materialization) { + if (analysis.anchors.size() != materialization.anchors.size()) + return false; + for (auto [analysisAnchor, materializationAnchor] : + llvm::zip_equal(analysis.anchors, materialization.anchors)) + if (analysisAnchor.kind != materializationAnchor.kind || + analysisAnchor.materializable != materializationAnchor.materializable) + return false; + return true; +} + static LogicalResult appendJSONLine(llvm::StringRef path, llvm::StringRef json) { if (path.empty()) @@ -127,6 +144,24 @@ struct SelectSimdSimtCostModelPass clearPreviousSelection(module); const bool autoMode = mode.getValue() == "auto"; + // Selection may inspect a transformed analysis view while materializing + // the chosen route on the route-neutral module owned by this pass. This + // is how AutoBlockify V1 dispatch/loop cost becomes visible without + // forcing an all-SIMD result to retain SIMT scheduling IR. + OwningOpRef parsedAnalysisModule; + ModuleOp analysisModule = module; + if (!analysisModulePath.getValue().empty()) { + parsedAnalysisModule = parseSourceFile( + analysisModulePath.getValue(), module.getContext()); + if (!parsedAnalysisModule) { + module.emitError("failed to parse SIMD/SIMT analysis module: ") + << analysisModulePath.getValue(); + signalPassFailure(); + return; + } + analysisModule = *parsedAnalysisModule; + } + SimdSimtCostModelOptions options; options.profilePath = profilePath.getValue(); options.actualTarget = actualTarget.getValue(); @@ -138,10 +173,20 @@ struct SelectSimdSimtCostModelPass wholeKernelSuperblockMaterializable.getValue(); options.scopeSuperblockMaterializable = scopeSuperblockMaterializable.getValue(); + options.logicalProgramCountHint = + std::max(0, logicalProgramCountHint.getValue()); + if (auto capability = + llvm::json::parse(routeTransformCapabilityJSON.getValue())) + if (auto *object = capability->getAsObject()) + if (auto count = object->getInteger("physical_vector_core_count_hint")) + options.physicalVectorCoreCountHint = std::max(0, *count); SimtAnchorPlan anchorPlan = buildMixedSimtAnchorPlan(module, options.compileOn91095); - auto reportOr = analyzeSimdSimtCandidates(module, anchorPlan, options); + SimtAnchorPlan analysisAnchorPlan = + buildMixedSimtAnchorPlan(analysisModule, options.compileOn91095); + auto reportOr = + analyzeSimdSimtCandidates(analysisModule, analysisAnchorPlan, options); if (!reportOr) { module.emitError("C++ SIMD/SIMT cost model failed: ") << llvm::toString(reportOr.takeError()); @@ -150,7 +195,10 @@ struct SelectSimdSimtCostModelPass } SimdSimtCostReport report = std::move(*reportOr); - std::string recommended = stringifySimdSimtCandidate(report.decision).str(); + std::string recommended = + report.stageModel.applied + ? stringifySimdSimtCandidate(report.decision).str() + : kBackendDefault.str(); std::string effective = kBackendDefault.str(); std::string selectionSource = "backend_default"; std::string applicationReason; @@ -169,12 +217,18 @@ struct SelectSimdSimtCostModelPass report.stageModel.mixed.routeSuperblockFactor; } - bool actionSupported = true; + bool actionSupported = report.stageModel.applied; + if (!report.stageModel.applied) + applicationReason = "stage_model_not_applicable"; bool hasExplicitScope = containsExplicitVectorScope(module); if (recommended == kMixedSimdSimt) { if (hasExplicitScope) { actionSupported = false; applicationReason = "explicit_scope_present"; + } else if (!anchorPlansHaveCompatibleIndices(analysisAnchorPlan, + anchorPlan)) { + actionSupported = false; + applicationReason = "analysis_materialization_anchor_mismatch"; } else { selectedMixedAnchorPlan = buildSelectedMixedAnchorPlan(report.stageModel, anchorPlan); @@ -199,8 +253,12 @@ struct SelectSimdSimtCostModelPass actionSupported = false; applicationReason = "explicit_scope_present"; } - if (selectedSuperblockFactor > 1 && - selectedSuperblockFactor * options.numWarps > 64) { + // Both whole-kernel and mixed-kernel V1 schedules launch + // numWarps * factor logical warp groups. Treating a mixed factor as the + // total warp count understated the resource limit by numWarps. + const int64_t selectedWarpCount = + selectedSuperblockFactor * options.numWarps; + if (selectedSuperblockFactor > 1 && selectedWarpCount > 64) { actionSupported = false; applicationReason = "superblock_warp_limit_exceeded"; } @@ -215,7 +273,7 @@ struct SelectSimdSimtCostModelPass effective = recommended; selectionSource = "cpp_cost_model"; applicationReason = "minimum_cost_candidate"; - } else if (!autoMode) { + } else if (!autoMode && report.stageModel.applied) { applicationReason = "report_mode"; } else if (applicationReason.empty()) { applicationReason = "candidate_not_materializable"; @@ -227,33 +285,70 @@ struct SelectSimdSimtCostModelPass module->setAttr(kEffectiveExecutionAttr, builder.getStringAttr(effective)); module->setAttr(kSelectionSourceAttr, builder.getStringAttr(selectionSource)); - module->setAttr(kAllSimdScoreAttr, - builder.getF64FloatAttr(report.candidateCosts.allSimd)); - module->setAttr(kAllSimtScoreAttr, - builder.getF64FloatAttr(report.candidateCosts.allSimtOnly)); - module->setAttr(kMixedScoreAttr, builder.getF64FloatAttr( - report.candidateCosts.mixedSimdSimt)); + if (report.stageModel.applied) { + module->setAttr(kAllSimdScoreAttr, + builder.getF64FloatAttr(report.candidateCosts.allSimd)); + module->setAttr( + kAllSimtScoreAttr, + builder.getF64FloatAttr(report.candidateCosts.allSimtOnly)); + module->setAttr( + kMixedScoreAttr, + builder.getF64FloatAttr(report.candidateCosts.mixedSimdSimt)); + } else { + module->removeAttr(kAllSimdScoreAttr); + module->removeAttr(kAllSimtScoreAttr); + module->removeAttr(kMixedScoreAttr); + } module->setAttr(kSuperblockFactorAttr, builder.getI64IntegerAttr(selectedSuperblockFactor)); // Selector and Materializer consume the same immutable anchor plan in one // pass invocation. No per-operation marker is persisted in TTIR. if (effective == kMixedSimdSimt && - failed(materializeSimtAnchorPlan(module, selectedMixedAnchorPlan))) { + failed(materializeSimtAnchorPlan(module, selectedMixedAnchorPlan, + selectedSuperblockFactor))) { signalPassFailure(); return; } llvm::json::Object reportJSON = report.toJSON(); + if (!report.stageModel.applied) { + // No Legacy/aggregate fallback exists. Do not publish + // default-initialized candidate scores as if the Stage model had + // evaluated this kernel. + reportJSON.erase("candidate_costs"); + reportJSON["decision_kind"] = kBackendDefault; + reportJSON["selectable_candidates"] = llvm::json::Array(); + } reportJSON["mode"] = mode.getValue(); reportJSON["recommended_decision_kind"] = recommended; reportJSON["effective_decision_kind"] = effective; reportJSON["selection_source"] = selectionSource; reportJSON["application_reason"] = applicationReason; reportJSON["action_supported"] = actionSupported; + reportJSON["analysis_ir_source"] = analysisModule == module + ? "route_neutral_ttir" + : "post_auto_blockify_v1_ttir"; + if (auto capability = + llvm::json::parse(routeTransformCapabilityJSON.getValue())) { + reportJSON["route_transform_capability"] = std::move(*capability); + } else { + module.emitError("invalid route-transform-capability-json"); + signalPassFailure(); + return; + } reportJSON["materialized_simt_anchor_count"] = static_cast(mixedAnchors.size()); reportJSON["selected_superblock_factor"] = selectedSuperblockFactor; + reportJSON["logical_program_count_hint"] = options.logicalProgramCountHint; + if (options.logicalProgramCountHint > 0) { + reportJSON["effective_runtime_factor"] = std::min( + selectedSuperblockFactor, options.logicalProgramCountHint); + reportJSON["full_group_count"] = + options.logicalProgramCountHint / selectedSuperblockFactor; + reportJSON["tail_count"] = + options.logicalProgramCountHint % selectedSuperblockFactor; + } std::string json = llvm::formatv("{0}", llvm::json::Value(std::move(reportJSON))).str(); module->setAttr(kReportJSONAttr, builder.getStringAttr(json)); diff --git a/third_party/ascend/costmodel/profiles/README.md b/third_party/ascend/costmodel/profiles/README.md index 12611dd332..edd04469c9 100644 --- a/third_party/ascend/costmodel/profiles/README.md +++ b/third_party/ascend/costmodel/profiles/README.md @@ -4,7 +4,7 @@ This directory is the canonical source for target-specific Cost Model data. - `microbench/` contains model-neutral hardware measurements shared by the absolute/autotune model and the SIMD/SIMT Route Model. -- `simd_simt/` contains Route Model policy, calibration, schema, and DES +- `simd_simt/` contains Route Model policy, calibration, and schema feedback data. Python packaging copies these files to diff --git a/third_party/ascend/costmodel/profiles/simd_simt/david_v100_des_feedback_v1.json b/third_party/ascend/costmodel/profiles/simd_simt/david_v100_des_feedback_v1.json deleted file mode 100644 index eb4428b871..0000000000 --- a/third_party/ascend/costmodel/profiles/simd_simt/david_v100_des_feedback_v1.json +++ /dev/null @@ -1,18 +0,0 @@ -{ - "schema_version": 1, - "kind": "simd_simt_calibration_feedback", - "profile_version": "david-v100-des-feedback-20260727-v1", - "feedback_version": "david-v100-des-feedback-20260727-v1", - "target": "Ascend950PR/dav-c310", - "base_profile_version": "david-v100-simd-simt-20260727-v3", - "rules": [], - "validation": { - "event_validated": false, - "maximum_fallback_ratio": 0.1 - }, - "provenance": { - "producer": "HIVMAnalysisReport JSON", - "policy": "Only enable a rule after op coverage and source mapping pass their gates", - "current_status": "Interface enabled; no correction is applied because the current gather-dot DES run has 1672/1727 scalar-or-unknown fallback operations and FBGEMM three-candidate coverage is incomplete" - } -} diff --git a/third_party/ascend/costmodel/profiles/simd_simt/david_v100_simd_simt_v1.json b/third_party/ascend/costmodel/profiles/simd_simt/david_v100_simd_simt_v1.json index 3d5ba5e707..67155babd3 100644 --- a/third_party/ascend/costmodel/profiles/simd_simt/david_v100_simd_simt_v1.json +++ b/third_party/ascend/costmodel/profiles/simd_simt/david_v100_simd_simt_v1.json @@ -1,7 +1,7 @@ { "$schema": "./simd_simt_profile_schema.json", "schema_version": 10, - "profile_version": "david-v100-simd-simt-20260820-v17", + "profile_version": "david-v100-simd-simt-20260824-v19", "target": "Ascend950PR/dav-c310", "description": "Stage Route Model hardware profile. Stage partitions and workloads are generated online from transformed TTIR; this file contains no workload-specific Stage template or route multiplier.", "microbenchmark_profile": "../microbench/ascend_davidv100_v1.json", @@ -130,7 +130,7 @@ "description": "All-SIMD dot setup and throughput seed used for every dot operation; current calibrated admission is narrower and covers only tiny irregular dot.", "startup_system_cycles": 128.0, "flops_per_system_cycle": 4096.0, - "source": "AscendModel cube startup/throughput seed; DES feedback is expected to refine scheduled lowering", + "source": "AscendModel cube startup/throughput seed; CaModel and profiler data refine scheduled lowering", "confidence": "low" }, "stage_resources": { @@ -290,12 +290,12 @@ "description": "Independent logical programs may continue hiding latency through F4. Loop-carried live state is pressure-free only through F2 and is charged above that point. Keeping the two limits separate lets straight-line FBGEMM-like work benefit from F4 without making recurrence-heavy solve_tril pressure-free." }, "scope_handoff": { - "fixed_directional_system_cycles": 0.0, + "fixed_directional_system_cycles": 668.0, "simd_ub_load_bytes_per_system_cycle": 512.0, "simd_ub_store_bytes_per_system_cycle": 256.0, "simt_ub_load_bytes_per_thread_per_system_cycle": 4.0, "simt_ub_store_bytes_per_thread_per_system_cycle": 4.0, - "description": "Only exact tensor values captured by or returned from scope.scope are charged. A value crosses the distinct SIMD/SIMT register files through UB: producer store plus consumer load. SIMD rates reflect the 256-byte vector width with two load and one store issue slots; SIMT rates are explicitly per active thread and are aggregated over one logical warp." + "description": "Only exact tensor values captured by or returned from scope.scope are charged. A value crosses the distinct SIMD/SIMT register files through UB: producer store plus consumer load. SIMD rates reflect the 256-byte vector width with two load and one store issue slots; SIMT rates are explicitly per active thread and are aggregated over one logical warp. The provisional 668-cycle directional fixed term is half of the 1335-cycle scope-pair residual obtained from the A5 local-scope sweep. One scope-local SuperBlock group pays this fixed pair once, so the per-logical-program route score amortizes only the fixed term by factor; UB bytes and Stage work remain unscaled. This is a route-independent calibration seed, not a solve_tril name check, and must later be replaced by directional transition microbenchmarks." }, "source": "Ascend950 SIMT issue limits plus provisional conservative dependency/reconvergence, SuperBlock state-pressure, and scope hand-off rates; replace numeric seeds only through versioned CaModel/microbenchmark evidence." }, diff --git a/third_party/ascend/include/CMakeLists.txt b/third_party/ascend/include/CMakeLists.txt index 9be89518ac..05abebd3c3 100644 --- a/third_party/ascend/include/CMakeLists.txt +++ b/third_party/ascend/include/CMakeLists.txt @@ -3,7 +3,8 @@ add_subdirectory(TritonToAnnotation) add_subdirectory(TritonToHFusion) add_subdirectory(TritonToHIVM) add_subdirectory(TritonToGraph) -add_subdirectory(TritonToLinalg) +add_subdirectory(TritonToLinalg) +add_subdirectory(TTIRLayoutMerge) add_subdirectory(TritonControlFlowOpt) add_subdirectory(Utils) add_subdirectory(DiscreteMaskAccessConversion) diff --git a/third_party/ascend/include/TTIRLayoutMerge/CMakeLists.txt b/third_party/ascend/include/TTIRLayoutMerge/CMakeLists.txt new file mode 100644 index 0000000000..983601b5c9 --- /dev/null +++ b/third_party/ascend/include/TTIRLayoutMerge/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name TTIRLayoutMerge) +add_public_tablegen_target(TTIRLayoutMergePassIncGen) diff --git a/third_party/ascend/include/TTIRLayoutMerge/Passes.h b/third_party/ascend/include/TTIRLayoutMerge/Passes.h new file mode 100644 index 0000000000..2c97730550 --- /dev/null +++ b/third_party/ascend/include/TTIRLayoutMerge/Passes.h @@ -0,0 +1,21 @@ +//===- Passes.h - TTIR layout-merge passes -------------------*- C++ -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef TRITON_ASCEND_TTIR_LAYOUT_MERGE_PASSES_H +#define TRITON_ASCEND_TTIR_LAYOUT_MERGE_PASSES_H + +#include "TTIRLayoutMergePass.h" + +namespace mlir::triton { + +#define GEN_PASS_REGISTRATION +#include "ascend/include/TTIRLayoutMerge/Passes.h.inc" + +} // namespace mlir::triton + +#endif // TRITON_ASCEND_TTIR_LAYOUT_MERGE_PASSES_H diff --git a/third_party/ascend/include/TTIRLayoutMerge/Passes.td b/third_party/ascend/include/TTIRLayoutMerge/Passes.td new file mode 100644 index 0000000000..4dc56d8b0b --- /dev/null +++ b/third_party/ascend/include/TTIRLayoutMerge/Passes.td @@ -0,0 +1,18 @@ +#ifndef TTIR_LAYOUT_MERGE_PASSES +#define TTIR_LAYOUT_MERGE_PASSES + +include "mlir/Pass/PassBase.td" + +def TTIRLayoutMerge : Pass<"ttir-layout-merge", "mlir::ModuleOp"> { + let summary = "Materialize TTIR layout merges before route cost analysis"; + let description = [{ + Applies implicit-permute normalization, strided-axis coalescing, + contiguous tile-chunk coalescing, and independent row coalescing while + the IR is still TTIR. The SIMD/SIMT Route Model is scheduled after this + pass so it observes the actual merged tensor shapes and memory-access + structure. + }]; + let constructor = "triton::createTTIRLayoutMergePass()"; +} + +#endif // TTIR_LAYOUT_MERGE_PASSES diff --git a/third_party/ascend/include/TritonToLinalg/RowCoalescing.h b/third_party/ascend/include/TTIRLayoutMerge/RowCoalescing.h similarity index 92% rename from third_party/ascend/include/TritonToLinalg/RowCoalescing.h rename to third_party/ascend/include/TTIRLayoutMerge/RowCoalescing.h index 23f30221a8..f1c5cca4fb 100644 --- a/third_party/ascend/include/TritonToLinalg/RowCoalescing.h +++ b/third_party/ascend/include/TTIRLayoutMerge/RowCoalescing.h @@ -20,8 +20,8 @@ * THE SOFTWARE. */ -#ifndef TRITON_ASCEND_ROW_COALESCING_H -#define TRITON_ASCEND_ROW_COALESCING_H +#ifndef TRITON_ASCEND_TTIR_LAYOUT_MERGE_ROW_COALESCING_H +#define TRITON_ASCEND_TTIR_LAYOUT_MERGE_ROW_COALESCING_H #include "mlir/Pass/Pass.h" #include "triton/Dialect/Triton/IR/Dialect.h" @@ -52,4 +52,4 @@ std::unique_ptr> createRowCoalescingPass(); } // namespace RowCoalescing -#endif // TRITON_ASCEND_ROW_COALESCING_H +#endif // TRITON_ASCEND_TTIR_LAYOUT_MERGE_ROW_COALESCING_H diff --git a/third_party/ascend/include/TritonToLinalg/TTIRLayoutMergePass.h b/third_party/ascend/include/TTIRLayoutMerge/TTIRLayoutMergePass.h similarity index 93% rename from third_party/ascend/include/TritonToLinalg/TTIRLayoutMergePass.h rename to third_party/ascend/include/TTIRLayoutMerge/TTIRLayoutMergePass.h index 31c9df3b53..7f3652bb17 100644 --- a/third_party/ascend/include/TritonToLinalg/TTIRLayoutMergePass.h +++ b/third_party/ascend/include/TTIRLayoutMerge/TTIRLayoutMergePass.h @@ -13,7 +13,7 @@ #include "mlir/Pass/Pass.h" #define GEN_PASS_DEF_TTIRLAYOUTMERGE -#include "ascend/include/TritonToLinalg/Passes.h.inc" +#include "ascend/include/TTIRLayoutMerge/Passes.h.inc" namespace mlir::triton { diff --git a/third_party/ascend/include/TritonToLinalg/Passes.h b/third_party/ascend/include/TritonToLinalg/Passes.h index dbe6d67f15..fe606b0700 100644 --- a/third_party/ascend/include/TritonToLinalg/Passes.h +++ b/third_party/ascend/include/TritonToLinalg/Passes.h @@ -24,7 +24,6 @@ #define TRITON_ADAPTER_TRITON_TO_LINALG_CONVERSION_PASSES_H #include "MarkTensorKindPass.h" -#include "TTIRLayoutMergePass.h" #include "TritonToLinalgPass.h" namespace mlir::triton { diff --git a/third_party/ascend/include/TritonToLinalg/Passes.td b/third_party/ascend/include/TritonToLinalg/Passes.td index 31d027c47f..729af06899 100644 --- a/third_party/ascend/include/TritonToLinalg/Passes.td +++ b/third_party/ascend/include/TritonToLinalg/Passes.td @@ -33,16 +33,4 @@ def MarkTensorKind : Pass<"mark-tensor-kind", "mlir::ModuleOp"> { let constructor = "triton::createMarkTensorKindPass()"; } -def TTIRLayoutMerge : Pass<"ttir-layout-merge", "mlir::ModuleOp"> { - let summary = "Materialize TTIR layout merges before route cost analysis"; - let description = [{ - Applies implicit-permute normalization, strided-axis coalescing, - contiguous tile-chunk coalescing, and independent row coalescing while - the IR is still TTIR. The SIMD / - SIMT Route Model is scheduled after this pass so it observes the actual - merged tensor shapes and memory-access structure. - }]; - let constructor = "triton::createTTIRLayoutMergePass()"; -} - #endif // TRITON_TO_LINALG_CONVERSION_PASSES diff --git a/third_party/ascend/lib/CMakeLists.txt b/third_party/ascend/lib/CMakeLists.txt index 1c2f67ad11..df0df2a3be 100644 --- a/third_party/ascend/lib/CMakeLists.txt +++ b/third_party/ascend/lib/CMakeLists.txt @@ -6,6 +6,7 @@ add_subdirectory(TritonToHFusion) add_subdirectory(TritonToHIVM) add_subdirectory(TritonToGraph) add_subdirectory(TritonToLinalg) +add_subdirectory(TTIRLayoutMerge) add_subdirectory(Utils) add_subdirectory(TritonControlFlowOpt) add_subdirectory(DiscreteMaskAccessConversion) @@ -22,6 +23,7 @@ if(TRITON_ENABLE_COVERAGE_HITEST) TritonToHFusion TritonToHIVM TritonToLinalg + TTIRLayoutMerge TritonControlFlowOpt TritonToLLVM TritonToStructured diff --git a/third_party/ascend/lib/DiscreteMaskAccessConversion/DiscreteMaskAccessConversionPass.cpp b/third_party/ascend/lib/DiscreteMaskAccessConversion/DiscreteMaskAccessConversionPass.cpp index b4e448b0e9..a5d2b6c6b3 100644 --- a/third_party/ascend/lib/DiscreteMaskAccessConversion/DiscreteMaskAccessConversionPass.cpp +++ b/third_party/ascend/lib/DiscreteMaskAccessConversion/DiscreteMaskAccessConversionPass.cpp @@ -20,7 +20,7 @@ * THE SOFTWARE. */ -#include "AscendModel/RouteModel/SimtSelection.h" +#include "AscendModel/Transforms/SimtSelection.h" #include "TritonToUnstructure/IndirectAtomicUtils.h" #include "Utils/Utils.h" #include "ascend/include/DiscreteMaskAccessConversion/Passes.h" diff --git a/third_party/ascend/lib/TTIRLayoutMerge/CMakeLists.txt b/third_party/ascend/lib/TTIRLayoutMerge/CMakeLists.txt new file mode 100644 index 0000000000..7725d33eb5 --- /dev/null +++ b/third_party/ascend/lib/TTIRLayoutMerge/CMakeLists.txt @@ -0,0 +1,18 @@ +add_triton_library(TTIRLayoutMerge + RowCoalescing.cpp + TTIRLayoutMergePass.cpp + + DEPENDS + TTIRLayoutMergePassIncGen + + LINK_LIBS PUBLIC + TritonToLinalg + MLIRArithDialect + MLIRControlFlowDialect + MLIRIR + MLIRMathDialect + MLIRPass + MLIRSCFDialect + MLIRTransforms + TritonIR +) diff --git a/third_party/ascend/lib/TritonToLinalg/RowCoalescing.cpp b/third_party/ascend/lib/TTIRLayoutMerge/RowCoalescing.cpp similarity index 99% rename from third_party/ascend/lib/TritonToLinalg/RowCoalescing.cpp rename to third_party/ascend/lib/TTIRLayoutMerge/RowCoalescing.cpp index c4ee85cbd0..38e5557611 100644 --- a/third_party/ascend/lib/TritonToLinalg/RowCoalescing.cpp +++ b/third_party/ascend/lib/TTIRLayoutMerge/RowCoalescing.cpp @@ -20,7 +20,7 @@ * THE SOFTWARE. */ -#include "TritonToLinalg/RowCoalescing.h" +#include "TTIRLayoutMerge/RowCoalescing.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" @@ -141,7 +141,6 @@ static std::optional matchRowSeed(ModuleOp moduleOp) { return RowSeed{pid, axis, cmp.getRhs(), cmp, falseBlock}; } - return std::nullopt; } diff --git a/third_party/ascend/lib/TritonToLinalg/TTIRLayoutMergePass.cpp b/third_party/ascend/lib/TTIRLayoutMerge/TTIRLayoutMergePass.cpp similarity index 96% rename from third_party/ascend/lib/TritonToLinalg/TTIRLayoutMergePass.cpp rename to third_party/ascend/lib/TTIRLayoutMerge/TTIRLayoutMergePass.cpp index 03e10c177e..f39fd56b39 100644 --- a/third_party/ascend/lib/TritonToLinalg/TTIRLayoutMergePass.cpp +++ b/third_party/ascend/lib/TTIRLayoutMerge/TTIRLayoutMergePass.cpp @@ -6,10 +6,10 @@ // //===----------------------------------------------------------------------===// -#include "ascend/include/TritonToLinalg/TTIRLayoutMergePass.h" +#include "ascend/include/TTIRLayoutMerge/TTIRLayoutMergePass.h" +#include "ascend/include/TTIRLayoutMerge/RowCoalescing.h" #include "ascend/include/TritonToLinalg/ImplicitPermute.h" -#include "ascend/include/TritonToLinalg/RowCoalescing.h" #include "ascend/include/TritonToLinalg/StridedAxisCoalescing.h" #include "ascend/include/TritonToLinalg/TileChunkCoalescing.h" diff --git a/third_party/ascend/lib/TritonToLinalg/CMakeLists.txt b/third_party/ascend/lib/TritonToLinalg/CMakeLists.txt index 95beb13463..7709ea5f3b 100644 --- a/third_party/ascend/lib/TritonToLinalg/CMakeLists.txt +++ b/third_party/ascend/lib/TritonToLinalg/CMakeLists.txt @@ -12,10 +12,8 @@ add_triton_library(TritonToLinalg StridedLoadStoreRewrite.cpp StridedAxisCoalescing.cpp TileChunkCoalescing.cpp - RowCoalescing.cpp DescriptorConverter.cpp MarkTensorKindPass.cpp - TTIRLayoutMergePass.cpp DevicePrintOffsetRewrite.cpp DEPENDS diff --git a/third_party/ascend/lib/TritonToLinalg/StridedLoadStoreRewrite.cpp b/third_party/ascend/lib/TritonToLinalg/StridedLoadStoreRewrite.cpp index 45cf0b7c5a..2beee0f20a 100644 --- a/third_party/ascend/lib/TritonToLinalg/StridedLoadStoreRewrite.cpp +++ b/third_party/ascend/lib/TritonToLinalg/StridedLoadStoreRewrite.cpp @@ -21,7 +21,7 @@ */ #include "TritonToLinalg/StridedLoadStoreRewrite.h" -#include "AscendModel/RouteModel/SimtSelection.h" +#include "AscendModel/Transforms/SimtSelection.h" #include "TritonToLinalg/ImplicitPermute.h" #include "TritonToLinalg/MaskAnalysis.h" #include "TritonToStructured/PtrAnalysis.h" diff --git a/third_party/ascend/lib/TritonToLinalg/TritonOpConverter.cpp b/third_party/ascend/lib/TritonToLinalg/TritonOpConverter.cpp index f9836e43f6..635ccc2c8e 100644 --- a/third_party/ascend/lib/TritonToLinalg/TritonOpConverter.cpp +++ b/third_party/ascend/lib/TritonToLinalg/TritonOpConverter.cpp @@ -22,7 +22,7 @@ */ #include "ascend/include/TritonToLinalg/TritonOpConverter.h" -#include "AscendModel/RouteModel/SimtSelection.h" +#include "AscendModel/Transforms/SimtSelection.h" #include "ascend/include/TritonToLinalg/BlockPtrAnalysis.h" #include "ascend/include/TritonToLinalg/MaskAnalysis.h" #include "ascend/include/TritonToLinalg/TritonToLinalgPass.h" diff --git a/third_party/ascend/lib/TritonToLinalg/TritonToLinalgPass.cpp b/third_party/ascend/lib/TritonToLinalg/TritonToLinalgPass.cpp index 5669fff84d..13bdf8503e 100644 --- a/third_party/ascend/lib/TritonToLinalg/TritonToLinalgPass.cpp +++ b/third_party/ascend/lib/TritonToLinalg/TritonToLinalgPass.cpp @@ -23,7 +23,7 @@ #include -#include "AscendModel/RouteModel/SimtSelection.h" +#include "AscendModel/Transforms/SimtSelection.h" #include "TritonToLinalg/BlockPtrAnalysis.h" #include "ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.h" #include "ascend/include/TritonToLinalg/ArgMinMaxConverter.h" @@ -96,6 +96,36 @@ using namespace mlir; using namespace triton; +/// Translate the frontend Route Model contract into the native BiShengIR +/// execution-region contract. `vector_mode` is intentionally the only +/// attribute authored by TTIR producers, but the downstream mixed pipeline +/// outlines a scope only when it carries the typed HIVM VF attributes. +/// Keeping this bridge here makes hand-written and auto-materialized scopes +/// follow the same lowering path without coupling the Route Model to HIVM. +static void materializeSimtScopeExecutionContract(ModuleOp module) { + module.walk([&](scope::ScopeOp scopeOp) { + auto mode = mlir::ascend::simt_selection::getVectorMode(scopeOp); + if (!mode || mode.getValue() != "simt") + return; + + OpBuilder builder(scopeOp); + // The Triton frontend historically emitted a discardable attribute named + // `noinline`. BiShengIR's scope dialect instead models `no_inline` as an + // inherent property, and InlineScope only consults that property. Make + // the boundary explicit here so the marked region survives until + // OutlineScope instead of being silently flattened back into SIMD code. + scopeOp->removeAttr("noinline"); + scopeOp.setNoInline(true); + scopeOp->setAttr("outline", builder.getUnitAttr()); + scopeOp->setAttr(hivm::TFuncCoreTypeAttr::name, + hivm::TFuncCoreTypeAttr::get(module.getContext(), + hivm::TFuncCoreType::AIV)); + scopeOp->setAttr( + hivm::VFModeAttr::name, + hivm::VFModeAttr::get(module.getContext(), hivm::VFMode::SIMT)); + }); +} + int nd2nzFlag = 0; bool compileOn91095Flag = false; bool existDotFlag = false; @@ -985,6 +1015,7 @@ void TritonToLinalgPass::runOnOperation() { compileModeFlag = triton::ascend::parseCompileMode(this->compileMode); auto moduleOp = getOperation(); + materializeSimtScopeExecutionContract(moduleOp); // Check if the kernel contains tl.dot. Without tl.dot, // the kernel would be pure AIV kernel. diff --git a/third_party/ascend/lib/TritonToUnstructure/UnstructureConversionPass.cpp b/third_party/ascend/lib/TritonToUnstructure/UnstructureConversionPass.cpp index c2ffbecdc4..392d863fe5 100644 --- a/third_party/ascend/lib/TritonToUnstructure/UnstructureConversionPass.cpp +++ b/third_party/ascend/lib/TritonToUnstructure/UnstructureConversionPass.cpp @@ -21,7 +21,7 @@ */ #include "TritonToUnstructure/UnstructureConversionPass.h" -#include "AscendModel/RouteModel/SimtSelection.h" +#include "AscendModel/Transforms/SimtSelection.h" #include "TritonToLinalg/MaskAnalysis.h" #include "TritonToStructured/CannonicalizerConverter.h" #include "TritonToUnstructure/IndirectAtomicUtils.h" diff --git a/third_party/ascend/triton_ascend.cc b/third_party/ascend/triton_ascend.cc index e547fea19a..9ee19d45e7 100644 --- a/third_party/ascend/triton_ascend.cc +++ b/third_party/ascend/triton_ascend.cc @@ -13,6 +13,7 @@ #include "ascend/include/AutoBlockifyV1/Passes.h" #include "ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.h" #include "ascend/include/DiscreteMaskAccessConversion/Passes.h" +#include "ascend/include/TTIRLayoutMerge/Passes.h" #include "ascend/include/TritonControlFlowOpt/Passes.h" #include "ascend/include/TritonToAnnotation/Passes.h" #include "ascend/include/TritonToHFusion/Passes.h" @@ -100,7 +101,10 @@ void init_triton_ascend_passes_ttir(py::module &&m) { const std::string &profilePath, const std::string &actualTarget, int64_t numWarps, bool compileOn91095, bool wholeKernelSuperblockMaterializable, - bool scopeSuperblockMaterializable, const std::string &reportFile) { + bool scopeSuperblockMaterializable, int64_t logicalProgramCountHint, + const std::string &analysisModulePath, + const std::string &routeTransformCapabilityJSON, + const std::string &reportFile) { mlir::ascend::SelectSimdSimtCostModelPassOptions opts; opts.mode = mode; opts.profilePath = profilePath; @@ -110,6 +114,9 @@ void init_triton_ascend_passes_ttir(py::module &&m) { opts.wholeKernelSuperblockMaterializable = wholeKernelSuperblockMaterializable; opts.scopeSuperblockMaterializable = scopeSuperblockMaterializable; + opts.logicalProgramCountHint = logicalProgramCountHint; + opts.analysisModulePath = analysisModulePath; + opts.routeTransformCapabilityJSON = routeTransformCapabilityJSON; opts.reportFile = reportFile; pm.addPass(mlir::ascend::createSelectSimdSimtCostModelPass(opts)); }, @@ -118,6 +125,9 @@ void init_triton_ascend_passes_ttir(py::module &&m) { py::arg("compile_on_910_95"), py::arg("whole_kernel_superblock_materializable") = false, py::arg("scope_superblock_materializable") = false, + py::arg("logical_program_count_hint") = 0, + py::arg("analysis_module_path") = "", + py::arg("route_transform_capability_json") = "{}", py::arg("report_file") = ""); m.def("add_materialize_simt_scopes", [](mlir::PassManager &pm) { @@ -453,10 +463,10 @@ runAscendCostModelInProcess(const std::string &mlirText, return os.str(); } -static std::string runHIVMCostModelInProcess( - const std::string &mlirText, const std::string &hardwareConfigPath, - const std::string &scheduler, const std::string &argBindings = "", - bool feedbackJSON = false) { +static std::string +runHIVMCostModelInProcess(const std::string &mlirText, + const std::string &hardwareConfigPath, + const std::string &argBindings = "") { std::string loadError; auto config = mlir::ascend::loadHardwareConfigForAnalysis(hardwareConfigPath, loadError); @@ -464,17 +474,6 @@ static std::string runHIVMCostModelInProcess( throw std::runtime_error("failed to load HIVM hardware config: " + loadError); - mlir::ascend::HIVMSchedulerMode mode; - if (scheduler.empty() || scheduler == "static") - mode = mlir::ascend::HIVMSchedulerMode::Static; - else if (scheduler == "des") - mode = mlir::ascend::HIVMSchedulerMode::DES; - else - throw std::runtime_error( - "invalid HIVM scheduler; expected 'static' or 'des'"); - if (feedbackJSON && mode != mlir::ascend::HIVMSchedulerMode::DES) - throw std::runtime_error("HIVM feedback JSON requires the 'des' scheduler"); - int fd = -1; llvm::SmallString<128> path; if (auto ec = llvm::sys::fs::createTemporaryFile("triton-hivm-costmodel", @@ -486,22 +485,16 @@ static std::string runHIVMCostModelInProcess( os << mlirText; } - mlir::ascend::HIVMAnalyzer analyzer(*config, argBindings, mode); + mlir::ascend::HIVMAnalyzer analyzer(*config, argBindings); mlir::ascend::HIVMAnalysisReport report; std::string error; bool ok = analyzer.analyzeFile(path, report, error); llvm::sys::fs::remove(path); if (!ok) throw std::runtime_error("in-process HIVM analysis failed: " + error); - if (feedbackJSON) - report.sourcePath = ""; - std::string result; llvm::raw_string_ostream os(result); - if (feedbackJSON) - report.emitFeedbackJSON(os, *config); - else - report.print(os, *config); + report.print(os, *config); return os.str(); } #endif @@ -537,23 +530,13 @@ void init_triton_ascend(py::module &&m) { m.def( "run_hivm_costmodel_inproc", [](const std::string &mlirText, const std::string &hardwareConfigPath, - const std::string &scheduler) { - py::gil_scoped_release release; - return runHIVMCostModelInProcess(mlirText, hardwareConfigPath, - scheduler); - }, - py::arg("mlir_text"), py::arg("hardware_config"), - py::arg("scheduler") = "des"); - m.def( - "run_hivm_costmodel_feedback_inproc", - [](const std::string &mlirText, const std::string &hardwareConfigPath, - const std::string &scheduler, const std::string &argBindings) { + const std::string &argBindings) { py::gil_scoped_release release; return runHIVMCostModelInProcess(mlirText, hardwareConfigPath, - scheduler, argBindings, true); + argBindings); }, py::arg("mlir_text"), py::arg("hardware_config"), - py::arg("scheduler") = "des", py::arg("arg_bindings") = ""); + py::arg("arg_bindings") = ""); #else m.def( "run_costmodel_inproc", @@ -571,17 +554,7 @@ void init_triton_ascend(py::module &&m) { return std::string(); }, py::arg("mlir_text"), py::arg("hardware_config"), - py::arg("scheduler") = "des"); - m.def( - "run_hivm_costmodel_feedback_inproc", - [](const std::string &, const std::string &, const std::string &, - const std::string &) { - throw std::runtime_error( - "in-process costmodel bridge is not enabled in this build"); - return std::string(); - }, - py::arg("mlir_text"), py::arg("hardware_config"), - py::arg("scheduler") = "des", py::arg("arg_bindings") = ""); + py::arg("arg_bindings") = ""); #endif // Initialize ascend IR bindings (ascendnpu_ir_builder, scope/hivm dialects) diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/simt_scope.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/simt_scope.mlir index 58e7ffef2b..36773da28b 100644 --- a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/simt_scope.mlir +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/simt_scope.mlir @@ -4,6 +4,12 @@ // the SIMT-aware runtime launch path. // CHECK-LABEL: func.func @simt_scope // CHECK-SAME: parallel_mode = "mix_simd_simt" +// CHECK: scope.scope +// CHECK-SAME: hivm.func_core_type = #hivm.func_core_type +// CHECK-SAME: hivm.vf_mode = #hivm.vf_mode +// CHECK-SAME: no_inline +// CHECK-SAME: outline +// CHECK-SAME: vector_mode = "simt" tt.func public @simt_scope(%arg0: !tt.ptr) { scope.scope : () -> () { scope.return diff --git a/third_party/ascend/unittest/costmodel_ut/PassesTest.cpp b/third_party/ascend/unittest/costmodel_ut/PassesTest.cpp index 6777c8db2b..51b63bb1ae 100644 --- a/third_party/ascend/unittest/costmodel_ut/PassesTest.cpp +++ b/third_party/ascend/unittest/costmodel_ut/PassesTest.cpp @@ -1,8 +1,8 @@ #include "AscendModel/Transforms/Passes.h" +#include "AscendModel/Analysis/SimtAnchorAnalysis.h" #include "AscendModel/IR/AscendModelDialect.h" #include "AscendModel/RouteModel/SimdSimtCostModel.h" -#include "AscendModel/RouteModel/SimtAnchorAnalysis.h" -#include "AscendModel/RouteModel/SimtSelection.h" +#include "AscendModel/Transforms/SimtSelection.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Func/IR/FuncOps.h" @@ -15,7 +15,10 @@ #include "mlir/Pass/PassManager.h" #include "bishengir/Dialect/Scope/IR/Scope.h" +#include "llvm/ADT/SmallString.h" +#include "llvm/Support/FileSystem.h" #include "llvm/Support/JSON.h" +#include "llvm/Support/raw_ostream.h" #include @@ -206,7 +209,7 @@ module { auto plan = buildMixedSimtAnchorPlan(*module, /*compileOn91095=*/true); ASSERT_EQ(plan.anchors.size(), 1u); - EXPECT_EQ(plan.materializableCount(), 1); + EXPECT_EQ(plan.materializableRoots().size(), 1u); EXPECT_EQ(mlir::ascend::stringifySimtAnchorKind(plan.anchors[0].kind), "loaded_index_dependent_memory"); @@ -215,26 +218,7 @@ module { FAIL() << llvm::toString(features.takeError()); EXPECT_EQ(features->loadedIndexDependentMemoryOps, 1); - EXPECT_EQ(features->simtAnchors.loadedIndexDependentMemoryOps, 1); - EXPECT_EQ(features->simtAnchors.recognizedCount, 1); EXPECT_EQ(features->simtAnchors.count, 1); - ASSERT_EQ(features->simtAnchors.mechanismKinds.size(), 1u); - EXPECT_EQ(features->simtAnchors.mechanismKinds.front(), - "loaded_index_dependent_memory"); - - // The same predicate is produced once and consumed by tt.load. The legacy - // rank sum sees both uses, while the hardware-facing fields count one SSA - // value and its 64 predicate elements exactly once. - EXPECT_GT(features->maskRankSum, features->uniqueMaskRankSum); - EXPECT_EQ(features->uniqueMaskValues, 1); - EXPECT_EQ(features->uniqueMaskRankSum, 2); - EXPECT_EQ(features->predicateElements, 64); - EXPECT_EQ(features->simtAnchors.uniqueMaskValues, 1); - EXPECT_EQ(features->simtAnchors.predicateElements, 64); - - EXPECT_EQ(features->rowLocalReduceOps, 1); - EXPECT_EQ(features->maxReduceAxisExtent, 16); - EXPECT_EQ(features->weightedReduceAxisElements, 16); } TEST(CostModelPassesTest, @@ -267,17 +251,13 @@ module { )mlir"); ASSERT_TRUE(module); + // Loop dependency classification is owned by StageFeatureAnalysis; the + // kernel summary deliberately no longer duplicates it. auto plan = buildMixedSimtAnchorPlan(*module, /*compileOn91095=*/true); - auto features = analyzeSimdSimtFeatures(*module, plan); - if (!features) - FAIL() << llvm::toString(features.takeError()); - - EXPECT_EQ(features->pointerInductionDependencyCount, 1); - EXPECT_EQ(features->loopCarriedDataDependencyCount, 1); + EXPECT_TRUE(plan.anchors.empty()); } -TEST(CostModelPassesTest, - SimtAnchorAnalysisExtractsHistogramFactsAndLowerability) { +TEST(CostModelPassesTest, SimtAnchorAnalysisClassifiesHistogramLowerability) { mlir::MLIRContext context; context.allowUnregisteredDialects(); auto module = parseModule(context, R"mlir( @@ -295,30 +275,16 @@ module { ASSERT_EQ(plan.anchors.size(), 1u); const auto &anchor = plan.anchors.front(); EXPECT_EQ(anchor.kind, mlir::ascend::SimtAnchorKind::Histogram); - const auto *facts = std::get_if(&anchor.facts); - ASSERT_NE(facts, nullptr); - EXPECT_EQ(facts->inputElements, 64); - EXPECT_EQ(facts->numBins, 256); - EXPECT_EQ(facts->inputType, "i32"); - EXPECT_EQ(facts->resultType, "i32"); - - EXPECT_EQ(anchor.lowerability.allSimd, - mlir::ascend::CandidateLoweringStatus::Unsupported); - EXPECT_EQ(anchor.lowerability.allSimtOnly, - mlir::ascend::CandidateLoweringStatus::Unsupported); - EXPECT_EQ(anchor.lowerability.mixed, - mlir::ascend::CandidateLoweringStatus::Native); + EXPECT_FALSE(anchor.lowerability.allSimd); + EXPECT_FALSE(anchor.lowerability.allSimtOnly); + EXPECT_TRUE(anchor.lowerability.mixed); EXPECT_TRUE(anchor.materializable); - EXPECT_EQ(plan.kernelLowerability.allSimd, - mlir::ascend::CandidateLoweringStatus::Unsupported); - EXPECT_EQ(plan.kernelLowerability.allSimtOnly, - mlir::ascend::CandidateLoweringStatus::Unsupported); - EXPECT_EQ(plan.kernelLowerability.mixed, - mlir::ascend::CandidateLoweringStatus::Native); + EXPECT_FALSE(plan.kernelLowerability.allSimd); + EXPECT_FALSE(plan.kernelLowerability.allSimtOnly); + EXPECT_TRUE(plan.kernelLowerability.mixed); } -TEST(CostModelPassesTest, - SimtAnchorAnalysisExtractsPlainCumsumFactsAndLowerability) { +TEST(CostModelPassesTest, SimtAnchorAnalysisClassifiesPlainCumsumLowerability) { mlir::MLIRContext context; context.allowUnregisteredDialects(); auto module = parseModule(context, R"mlir( @@ -342,30 +308,17 @@ module { const auto &anchor = plan.anchors.front(); EXPECT_EQ(anchor.kind, mlir::ascend::SimtAnchorKind::PlainOneDimensionalCumsum); - const auto *facts = - std::get_if(&anchor.facts); - ASSERT_NE(facts, nullptr); - EXPECT_EQ(facts->axisExtent, 128); - EXPECT_EQ(facts->elementType, "f32"); - EXPECT_TRUE(facts->reverse); - - EXPECT_EQ(anchor.lowerability.allSimd, - mlir::ascend::CandidateLoweringStatus::AliasesMixed); - EXPECT_EQ(anchor.lowerability.allSimtOnly, - mlir::ascend::CandidateLoweringStatus::BackendConditional); - EXPECT_EQ(anchor.lowerability.mixed, - mlir::ascend::CandidateLoweringStatus::Native); + EXPECT_FALSE(anchor.lowerability.allSimd); + EXPECT_TRUE(anchor.lowerability.allSimtOnly); + EXPECT_TRUE(anchor.lowerability.mixed); EXPECT_TRUE(anchor.materializable); - EXPECT_EQ(plan.kernelLowerability.allSimd, - mlir::ascend::CandidateLoweringStatus::AliasesMixed); - EXPECT_EQ(plan.kernelLowerability.allSimtOnly, - mlir::ascend::CandidateLoweringStatus::BackendConditional); - EXPECT_EQ(plan.kernelLowerability.mixed, - mlir::ascend::CandidateLoweringStatus::Native); + EXPECT_FALSE(plan.kernelLowerability.allSimd); + EXPECT_TRUE(plan.kernelLowerability.allSimtOnly); + EXPECT_TRUE(plan.kernelLowerability.mixed); } TEST(CostModelPassesTest, - SimtAnchorAnalysisExtractsTensorAtomicFactsAndLowerability) { + SimtAnchorAnalysisClassifiesTensorAtomicLowerability) { mlir::MLIRContext context; context.allowUnregisteredDialects(); auto module = parseModule(context, R"mlir( @@ -393,34 +346,13 @@ module { ASSERT_EQ(plan.anchors.size(), 1u); const auto &anchor = plan.anchors.front(); EXPECT_EQ(anchor.kind, mlir::ascend::SimtAnchorKind::TensorAtomic); - const auto *facts = - std::get_if(&anchor.facts); - ASSERT_NE(facts, nullptr); - EXPECT_EQ(facts->updateElements, 64); - EXPECT_EQ(facts->addressRank, 1); - EXPECT_EQ(facts->valueType, "f32"); - EXPECT_EQ(facts->offsetType, "i64"); - EXPECT_EQ(facts->operation, "fadd"); - EXPECT_TRUE(facts->hasMask); - ASSERT_TRUE(facts->staticMaskActiveFraction.has_value()); - EXPECT_DOUBLE_EQ(*facts->staticMaskActiveFraction, 1.0); - EXPECT_TRUE(facts->resultUsed); - EXPECT_TRUE(facts->addressIsLaneVarying); - EXPECT_TRUE(facts->addressDependsOnLoadedIndex); - - EXPECT_EQ(anchor.lowerability.allSimd, - mlir::ascend::CandidateLoweringStatus::Native); - EXPECT_EQ(anchor.lowerability.allSimtOnly, - mlir::ascend::CandidateLoweringStatus::BackendConditional); - EXPECT_EQ(anchor.lowerability.mixed, - mlir::ascend::CandidateLoweringStatus::Native); + EXPECT_TRUE(anchor.lowerability.allSimd); + EXPECT_TRUE(anchor.lowerability.allSimtOnly); + EXPECT_TRUE(anchor.lowerability.mixed); EXPECT_TRUE(anchor.materializable); - EXPECT_EQ(plan.kernelLowerability.allSimd, - mlir::ascend::CandidateLoweringStatus::Native); - EXPECT_EQ(plan.kernelLowerability.allSimtOnly, - mlir::ascend::CandidateLoweringStatus::BackendConditional); - EXPECT_EQ(plan.kernelLowerability.mixed, - mlir::ascend::CandidateLoweringStatus::Native); + EXPECT_TRUE(plan.kernelLowerability.allSimd); + EXPECT_TRUE(plan.kernelLowerability.allSimtOnly); + EXPECT_TRUE(plan.kernelLowerability.mixed); } TEST(CostModelPassesTest, @@ -478,11 +410,10 @@ module { auto plan = buildMixedSimtAnchorPlan(*module, /*compileOn91095=*/true); ASSERT_EQ(plan.anchors.size(), 1u); - EXPECT_EQ(plan.materializableCount(), 1); + EXPECT_EQ(plan.materializableRoots().size(), 1u); for (const auto &anchor : plan.anchors) { EXPECT_EQ(anchor.kind, mlir::ascend::SimtAnchorKind::TriangularSolveLoop); - EXPECT_EQ(anchor.lowerability.mixed, - mlir::ascend::CandidateLoweringStatus::Native); + EXPECT_TRUE(anchor.lowerability.mixed); EXPECT_TRUE(anchor.materializable); // Both recurrence loops are one physical SIMT scope and therefore one // scored/materialized anchor, not two independent route decisions. @@ -490,43 +421,31 @@ module { EXPECT_EQ(anchor.scopeInsertionPoint, anchor.operation); EXPECT_EQ(anchor.scopeOperations.front()->getName().getStringRef(), "tt.make_range"); - const auto *facts = - std::get_if(&anchor.facts); - ASSERT_NE(facts, nullptr); - EXPECT_EQ(facts->blockRows, 16); - EXPECT_EQ(facts->blockColumns, 16); - EXPECT_EQ(facts->accumulatorType, "f32"); - EXPECT_EQ(facts->recurrenceStartRow, 2); - EXPECT_EQ(facts->recurrenceLoopCount, 2); - EXPECT_EQ(facts->denseDotTailOps, 0); - EXPECT_FALSE(facts->requiresCubeTailPartition); + ASSERT_TRUE(anchor.triangularSolve); + EXPECT_EQ(anchor.triangularSolve->blockRows, 16); + EXPECT_EQ(anchor.triangularSolve->blockColumns, 16); + EXPECT_EQ(anchor.triangularSolve->accumulatorType, "f32"); + EXPECT_EQ(anchor.triangularSolve->recurrenceStartRow, 2); + // Two recurrence loops, each with 14 body iterations. + EXPECT_EQ(anchor.triangularSolve->recurrenceLoopCount, 28); + EXPECT_EQ(anchor.triangularSolve->denseDotTailOps, 0); + EXPECT_FALSE(anchor.triangularSolve->requiresCubeTailPartition); } - EXPECT_EQ(plan.kernelLowerability.mixed, - mlir::ascend::CandidateLoweringStatus::Native); + EXPECT_TRUE(plan.kernelLowerability.mixed); auto features = analyzeSimdSimtFeatures(*module, plan); if (!features) FAIL() << llvm::toString(features.takeError()); EXPECT_EQ(features->simtAnchors.count, 1); - EXPECT_EQ(features->simtAnchors.staticLoopCount, 2); - EXPECT_EQ(features->simtAnchors.staticLoopTripCountSum, 28); - EXPECT_EQ(features->simtAnchors.modeledDynamicLoopCount, 2); - EXPECT_EQ(features->simtAnchors.modeledDynamicLoopTripCountSum, 28); - EXPECT_EQ(features->simtAnchors.reduceOps, 2); - EXPECT_EQ(features->simtAnchors.weightedOps.lookup("reduce"), 28); - EXPECT_EQ(features->simtAnchors.shuffleLaneSteps, 28672); - // 2 loop-weighted mask uses: 2 * 256 * 14, plus the 256-lane mask setup - // that is moved into the same scope. - EXPECT_EQ(features->simtAnchors.predicateLaneEvaluations, 7424); - EXPECT_TRUE(features->hasUnknownTripCount); - EXPECT_EQ(features->loopCarriedDataDependencyCount, 2); - EXPECT_EQ(features->pointerInductionDependencyCount, 0); - - ASSERT_TRUE(mlir::succeeded(materializeSimtAnchorPlan(*module, plan))); + ASSERT_TRUE(mlir::succeeded(materializeSimtAnchorPlan(*module, plan, 4))); Operation *scope = findFirstOp(*module, "scope.scope"); ASSERT_NE(scope, nullptr); EXPECT_EQ(scope->getAttrOfType("vector_mode").getValue(), "simt"); + EXPECT_EQ( + scope->getAttrOfType("ascend.scope_superblock.factor") + .getInt(), + 4); Operation *initialLoad = findFirstOp(*module, "tt.load"); ASSERT_NE(initialLoad, nullptr); EXPECT_NE(initialLoad->getParentOp(), scope); @@ -599,7 +518,7 @@ TEST(CostModelPassesTest, PerfReportPassAcceptsEstimatedPipeline) { createPipelineAnalysisPass(), createPerfReportPass())); } -TEST(CostModelPassesTest, SimdSimtAutoAlwaysScoresLegalCandidates) { +TEST(CostModelPassesTest, SimdSimtReturnsBackendDefaultOutsideStageDomain) { auto configureOptions = [](SelectSimdSimtCostModelPassOptions &options, llvm::StringRef mode) { options.mode = mode.str(); @@ -629,23 +548,19 @@ TEST(CostModelPassesTest, SimdSimtAutoAlwaysScoresLegalCandidates) { ASSERT_TRUE(autoEffective); ASSERT_TRUE(autoRecommended); ASSERT_TRUE(autoReport); - EXPECT_EQ(autoEffective.getValue(), "all_simd"); - EXPECT_EQ(autoRecommended.getValue(), "all_simd"); - EXPECT_TRUE((*autoModule)->hasAttr("ascend.simt_costmodel.all_simd_score")); + EXPECT_EQ(autoEffective.getValue(), "backend_default"); + EXPECT_EQ(autoRecommended.getValue(), "backend_default"); + EXPECT_FALSE((*autoModule)->hasAttr("ascend.simt_costmodel.all_simd_score")); auto autoJSON = llvm::json::parse(autoReport.getValue()); ASSERT_TRUE(static_cast(autoJSON)); auto *autoObject = autoJSON->getAsObject(); ASSERT_NE(autoObject, nullptr); - auto *autoCandidateCosts = autoObject->get("candidate_costs"); - auto *autoDecision = autoObject->get("decision_kind"); - ASSERT_NE(autoCandidateCosts, nullptr); - ASSERT_NE(autoDecision, nullptr); - EXPECT_NE(autoCandidateCosts->getAsObject(), nullptr); - ASSERT_TRUE(autoDecision->getAsString()); - EXPECT_EQ(*autoDecision->getAsString(), "all_simd"); + auto autoDecision = autoObject->getString("decision_kind"); + ASSERT_TRUE(autoDecision); + EXPECT_EQ(*autoDecision, "backend_default"); auto autoReason = autoObject->getString("application_reason"); ASSERT_TRUE(autoReason); - EXPECT_EQ(*autoReason, "minimum_cost_candidate"); + EXPECT_EQ(*autoReason, "stage_model_not_applicable"); mlir::MLIRContext reportContext; auto reportModule = parseModule(reportContext, kOutOfSimdSimtCoverageModule); @@ -664,17 +579,71 @@ TEST(CostModelPassesTest, SimdSimtAutoAlwaysScoresLegalCandidates) { ASSERT_TRUE(reportEffective); ASSERT_TRUE(reportJSONAttr); EXPECT_EQ(reportEffective.getValue(), "backend_default"); - EXPECT_TRUE((*reportModule)->hasAttr("ascend.simt_costmodel.all_simd_score")); + EXPECT_FALSE( + (*reportModule)->hasAttr("ascend.simt_costmodel.all_simd_score")); auto reportJSON = llvm::json::parse(reportJSONAttr.getValue()); ASSERT_TRUE(static_cast(reportJSON)); auto *reportObject = reportJSON->getAsObject(); ASSERT_NE(reportObject, nullptr); auto reportDecision = reportObject->getString("decision_kind"); ASSERT_TRUE(reportDecision); - EXPECT_EQ(*reportDecision, "all_simd"); + EXPECT_EQ(*reportDecision, "backend_default"); auto reportReason = reportObject->getString("application_reason"); ASSERT_TRUE(reportReason); - EXPECT_EQ(*reportReason, "report_mode"); + EXPECT_EQ(*reportReason, "stage_model_not_applicable"); +} + +TEST(CostModelPassesTest, SimdSimtSelectionUsesExternalAnalysisIR) { + mlir::MLIRContext context; + auto module = parseModule(context, kOutOfSimdSimtCoverageModule); + ASSERT_TRUE(module); + + llvm::SmallString<128> analysisPath; + int analysisFd = -1; + ASSERT_FALSE(llvm::sys::fs::createTemporaryFile( + "simd_simt_v1_analysis", "mlir", analysisFd, analysisPath)); + { + llvm::raw_fd_ostream analysisFile(analysisFd, true); + analysisFile << R"mlir( +module { + func.func @main(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> + attributes {ta.auto_blockify_v1} { + %0 = arith.addf %arg0, %arg1 : tensor<4xf32> + return %0 : tensor<4xf32> + } +} +)mlir"; + } + + SelectSimdSimtCostModelPassOptions options; + options.mode = "report"; + options.profilePath = TRITON_ASCEND_SIMD_SIMT_TEST_PROFILE_PATH; + options.actualTarget = "Ascend950PR_9579"; + options.numWarps = 4; + options.compileOn91095 = true; + options.analysisModulePath = analysisPath.str().str(); + const bool succeeded = + runPasses(*module, createSelectSimdSimtCostModelPass(options)); + llvm::sys::fs::remove(analysisPath); + ASSERT_TRUE(succeeded); + + auto reportAttr = + (*module)->getAttrOfType("ascend.simt_costmodel.report_json"); + ASSERT_TRUE(reportAttr); + auto report = llvm::json::parse(reportAttr.getValue()); + ASSERT_TRUE(static_cast(report)); + auto *object = report->getAsObject(); + ASSERT_NE(object, nullptr); + auto analysisSource = object->getString("analysis_ir_source"); + ASSERT_TRUE(analysisSource); + EXPECT_EQ(*analysisSource, "post_auto_blockify_v1_ttir"); + auto *features = object->getObject("features"); + ASSERT_NE(features, nullptr); + auto *postTransform = features->getObject("post_transform"); + ASSERT_NE(postTransform, nullptr); + auto v1Applied = postTransform->getBoolean("auto_blockify_v1_applied"); + ASSERT_TRUE(v1Applied); + EXPECT_TRUE(*v1Applied); } TEST(CostModelPassesTest, MaterializeSimtScopePreservesEscapingSSAResult) { @@ -734,6 +703,64 @@ module attributes { "ascend.simt_costmodel.scope_materialized")); } +TEST(CostModelPassesTest, SameStageAnchorsMaterializeAsOneCompoundScope) { + mlir::MLIRContext context; + auto module = parseModule(context, R"mlir( +module { + func.func @main(%arg0: i32, %arg1: i32) -> i32 { + %0 = arith.addi %arg0, %arg1 : i32 + %1 = arith.muli %0, %arg1 : i32 + %2 = arith.addi %1, %arg0 : i32 + %3 = arith.addi %0, %2 : i32 + return %3 : i32 + } +} +)mlir"); + ASSERT_TRUE(module); + + Operation *first = findFirstOp(*module, "arith.muli")->getPrevNode(); + Operation *middle = findFirstOp(*module, "arith.muli"); + Operation *second = middle->getNextNode(); + ASSERT_NE(first, nullptr); + ASSERT_NE(middle, nullptr); + ASSERT_NE(second, nullptr); + + mlir::ascend::SimtAnchorPlan plan; + for (Operation *operation : {first, second}) { + mlir::ascend::SimtAnchorDescriptor anchor; + anchor.operation = operation; + anchor.scopeOperations.push_back(operation); + anchor.scopeInsertionPoint = operation; + anchor.kind = mlir::ascend::SimtAnchorKind::LoadedIndexDependentMemory; + anchor.materializable = true; + plan.anchors.push_back(std::move(anchor)); + } + + auto merged = mlir::ascend::mergeSimtStageAnchors(plan, {0, 1}); + ASSERT_TRUE(merged); + ASSERT_EQ(merged->scopeOperations.size(), 3u); + EXPECT_EQ(merged->scopeOperations[0], first); + EXPECT_EQ(merged->scopeOperations[1], middle); + EXPECT_EQ(merged->scopeOperations[2], second); + + mlir::ascend::SimtAnchorPlan selected; + selected.anchors.push_back(std::move(*merged)); + ASSERT_TRUE(mlir::succeeded(materializeSimtAnchorPlan(*module, selected, 2))); + Operation *scope = findFirstOp(*module, "scope.scope"); + ASSERT_NE(scope, nullptr); + EXPECT_EQ( + scope->getAttrOfType("ascend.scope_superblock.factor") + .getInt(), + 2); + EXPECT_EQ(middle->getParentOp(), scope); + int64_t scopeCount = 0; + module->walk([&](Operation *operation) { + scopeCount += operation->getName().getStringRef() == "scope.scope"; + }); + EXPECT_EQ(scopeCount, 1); + EXPECT_TRUE(mlir::succeeded(mlir::verify(*module))); +} + TEST(CostModelPassesTest, NativeWholeBodySimtScopeDetectionAndInlining) { mlir::MLIRContext context; auto module = parseModule(context, R"mlir( diff --git a/third_party/ascend/unittest/costmodel_ut/SimdSimtCostModelTest.cpp b/third_party/ascend/unittest/costmodel_ut/SimdSimtCostModelTest.cpp index a3dee9442f..539997ebbc 100644 --- a/third_party/ascend/unittest/costmodel_ut/SimdSimtCostModelTest.cpp +++ b/third_party/ascend/unittest/costmodel_ut/SimdSimtCostModelTest.cpp @@ -1,6 +1,6 @@ #include "AscendModel/RouteModel/SimdSimtCostModel.h" +#include "AscendModel/Analysis/StagePartitioner.h" #include "AscendModel/RouteModel/StageCostModels.h" -#include "AscendModel/RouteModel/StagePartitioner.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Func/IR/FuncOps.h" @@ -9,19 +9,18 @@ #include -using mlir::ascend::estimateSimdSimtCandidates; using mlir::ascend::HardwareProfile; using mlir::ascend::LogicalPhase; +using mlir::ascend::LogicalPhaseCost; using mlir::ascend::LogicalStage; -using mlir::ascend::SimdSimtCandidateKind; -using mlir::ascend::SimdSimtCostModelOptions; +using mlir::ascend::LogicalStageCost; using mlir::ascend::SimdSimtFeatureSummary; using mlir::ascend::solveStageRoutes; using mlir::ascend::StageCostEvaluator; using mlir::ascend::StageCostModelKind; -using mlir::ascend::StageCostModelRegistry; using mlir::ascend::StageCostTable; using mlir::ascend::StageFeatureAnalysis; +using mlir::ascend::StageImplementationCost; using mlir::ascend::StageMode; using mlir::ascend::StageModeLegalityAnalysis; using mlir::ascend::StagePartition; @@ -35,190 +34,13 @@ using mlir::ascend::TriangularSolveFacts; namespace { -SimdSimtCostModelOptions options(unsigned numWarps) { - SimdSimtCostModelOptions result; - result.profilePath = TRITON_ASCEND_SIMD_SIMT_TEST_PROFILE_PATH; - result.actualTarget = "Ascend950PR_9579"; - result.numWarps = numWarps; - result.compileOn91095 = true; - return result; -} - -SimdSimtFeatureSummary gatherDotFeatures() { - SimdSimtFeatureSummary f; - f.loadOps = 3; - f.storeOps = 1; - f.dotOps = 1; - f.broadcastOps = 5; - f.expandDimsOps = 4; - f.splatOps = 9; - f.addPtrOps = 7; - f.arithOps = 9; - f.addOps = 2; - f.mulOps = 5; - f.scalarOps = 7; - f.maxTensorRank = 2; - f.maxTensorNumel = 256; - f.maxElementBits = 32; - f.pointerTensorOps = 11; - f.pointerUnstructuredDims = 18; - f.laneDependentPointerOps = 9; - f.loadedIndexDependentMemoryOps = 2; - f.vectorPtrSplatOps = 4; - f.loadBytes = 1088; - f.storeBytes = 1024; - f.loadWarpInstructions = 17; - f.storeWarpInstructions = 8; - f.dotFlops = 8192; - f.dotOutputElements = 256; - f.dotMNK.push_back({16, 16, 16}); - f.hasDot = true; - f.observedMixedKinds.push_back("conditional_indirect_memory"); - f.simtAnchors.count = 2; - f.simtAnchors.recognizedCount = 2; - f.simtAnchors.loadedIndexDependentMemoryOps = 2; - f.simtAnchors.mechanismKinds.push_back("loaded_index_dependent_memory"); - - f.weightedOps["add"] = 2; - f.weightedOps["mul"] = 5; - f.weightedOps["load"] = 3; - f.weightedOps["store"] = 1; - f.opElements["add"] = 32; - f.opElements["mul"] = 50; - f.opElements["load"] = 528; - f.opElements["store"] = 256; - return f; -} - -SimdSimtFeatureSummary fbgemmFeatures() { - SimdSimtFeatureSummary f; - f.loadOps = 7; - f.storeOps = 2; - f.reduceOps = 1; - f.broadcastOps = 6; - f.expandDimsOps = 6; - f.splatOps = 10; - f.addPtrOps = 12; - f.arithOps = 29; - f.mathOps = 2; - f.addOps = 1; - f.mulOps = 7; - f.divOps = 2; - f.maxOps = 2; - f.absOps = 2; - f.cmpOps = 1; - f.castOps = 2; - f.clampOps = 2; - f.scalarOps = 19; - f.maxTensorRank = 2; - f.maxTensorNumel = 64; - f.maxElementBits = 64; - f.maskTensorOps = 12; - f.maskRankSum = 25; - f.maskBroadcastOps = 4; - f.pointerTensorOps = 19; - f.pointerUnstructuredDims = 20; - f.laneDependentPointerOps = 10; - f.loadedIndexDependentMemoryOps = 3; - f.rowLocalReduceOps = 1; - f.scalarLoadOps = 2; - f.vectorPtrSplatOps = 6; - f.loadBytes = 4152; - f.storeBytes = 2064; - f.loadWarpInstructions = 37; - f.storeWarpInstructions = 17; - f.staticLoopCount = 2; - f.staticLoopTripCountSum = 16; - f.staticLoopTripCountMax = 8; - f.hasControlFlow = true; - f.observedMixedKinds.push_back("conditional_indirect_memory"); - f.simtAnchors.count = 3; - f.simtAnchors.recognizedCount = 3; - f.simtAnchors.loadedIndexDependentMemoryOps = 3; - f.simtAnchors.mechanismKinds.push_back("loaded_index_dependent_memory"); - - f.weightedOps["abs"] = 9; - f.weightedOps["add"] = 1; - f.weightedOps["cast"] = 2; - f.weightedOps["clamp"] = 9; - f.weightedOps["cmp"] = 1; - f.weightedOps["div"] = 2; - f.weightedOps["load"] = 21; - f.weightedOps["max"] = 16; - f.weightedOps["mul"] = 14; - f.weightedOps["reduce"] = 8; - f.weightedOps["store"] = 9; - f.opElements["abs"] = 516; - f.opElements["add"] = 4; - f.opElements["cast"] = 8; - f.opElements["clamp"] = 516; - f.opElements["cmp"] = 4; - f.opElements["div"] = 8; - f.opElements["load"] = 1038; - f.opElements["max"] = 40; - f.opElements["mul"] = 533; - f.opElements["reduce"] = 512; - f.opElements["store"] = 516; - return f; -} - -SimdSimtFeatureSummary outOfCoverageFeatures() { - SimdSimtFeatureSummary f = gatherDotFeatures(); - // One FLOP above the profile's tiny-dot coverage ceiling. Since dotFlops is - // non-zero, neither reduction-only coverage domain can admit this feature. - f.dotFlops = 16385; - return f; -} - -SimdSimtFeatureSummary rank1IndirectVectorReductionFeatures() { - SimdSimtFeatureSummary f; - f.reduceOps = 1; - f.maxTensorRank = 1; - f.maxTensorNumel = 256; - f.maxElementBits = 32; - f.rank1IndirectVectorReduce = true; - f.weightedOps["reduce"] = 8; - return f; -} - -SimdSimtFeatureSummary solveTrilBt16Features() { +SimdSimtFeatureSummary triangularBt16StageFeatures() { SimdSimtFeatureSummary f; f.reduceOps = 1; - f.maxTensorRank = 2; - f.maxTensorNumel = 256; - f.maxElementBits = 32; - f.maskRankSum = 20; - f.pointerTensorOps = 5; - f.laneDependentPointerOps = 2; - f.staticLoopCount = 1; - f.staticLoopTripCountSum = 1; - f.staticLoopTripCountMax = 1; - f.hasControlFlow = true; - f.weightedOps["reduce"] = 1; - return f; -} - -SimdSimtFeatureSummary triangularBt16StageFeatures() { - SimdSimtFeatureSummary f = solveTrilBt16Features(); f.loadOps = 1; f.storeOps = 1; - f.loadBytes = 16 * 16 * 4; - f.storeBytes = 16 * 16 * 4; - f.loadWarpInstructions = 8; - f.storeWarpInstructions = 8; - f.staticLoopCount = 1; - f.staticLoopTripCountSum = 14; f.staticLoopTripCountMax = 14; - f.predicateElements = 16 * 16 * 14; - f.predicateLaneEvaluations = 16 * 16 * 14; - f.shuffleLaneSteps = 16 * 16 * 4 * 14; f.simtAnchors.count = 1; - f.simtAnchors.recognizedCount = 1; - f.simtAnchors.reduceOps = 1; - f.simtAnchors.predicateElements = f.predicateElements; - f.simtAnchors.predicateLaneEvaluations = f.predicateLaneEvaluations; - f.simtAnchors.shuffleLaneSteps = f.shuffleLaneSteps; - f.simtAnchors.mechanismKinds.push_back("triangular_solve_loop"); TriangularSolveFacts triangular; triangular.blockRows = 16; triangular.blockColumns = 16; @@ -229,59 +51,6 @@ SimdSimtFeatureSummary triangularBt16StageFeatures() { return f; } -SimdSimtFeatureSummary triangularUnknownLoopFeatures() { - SimdSimtFeatureSummary f = solveTrilBt16Features(); - // BT64 has four sibling 16x16 recurrences. Their TTIR upper bounds remain - // dynamic, but the full-tile structural estimate is 14 trips per loop. - f.staticLoopCount = 4; - f.staticLoopTripCountSum = 56; - f.staticLoopTripCountMax = 14; - f.modeledDynamicLoopCount = 4; - f.modeledDynamicLoopTripCountSum = 56; - f.hasUnknownTripCount = true; - f.maskRankSum = 62; - f.loadOps = 4; - f.storeOps = 1; - f.loadBytes = 4 * 16 * 16 * 4; - f.storeBytes = 64 * 64 * 4; - f.loadWarpInstructions = 32; - f.storeWarpInstructions = 128; - f.weightedOps["reduce"] = 56; - f.shuffleLaneSteps = 57344; - f.predicateLaneEvaluations = 35840; - // The BT64 merge tail contains 16 dense 16x16x16 dots. They are outside - // the SIMT anchor and remain eligible for SIMD/Cube lowering. - f.dotOps = 16; - f.dotFlops = 16 * 2 * 16 * 16 * 16; - f.dotOutputElements = 16 * 16 * 16; - f.hasDot = true; - f.simtAnchors.count = 1; - f.simtAnchors.recognizedCount = 1; - f.simtAnchors.reduceOps = 4; - f.simtAnchors.maxTensorNumel = 256; - f.simtAnchors.maskRankSum = 62; - f.simtAnchors.staticLoopCount = 4; - f.simtAnchors.staticLoopTripCountSum = 56; - f.simtAnchors.modeledDynamicLoopCount = 4; - f.simtAnchors.modeledDynamicLoopTripCountSum = 56; - f.simtAnchors.weightedOps["reduce"] = 56; - f.simtAnchors.shuffleLaneSteps = 57344; - f.simtAnchors.predicateLaneEvaluations = 35840; - f.simtAnchors.mechanismKinds.push_back("triangular_solve_loop"); - f.simtAnchors.kernelLowerability.allSimtOnly = - mlir::ascend::CandidateLoweringStatus::BackendConditional; - TriangularSolveFacts triangular; - triangular.blockRows = 16; - triangular.blockColumns = 16; - triangular.accumulatorType = "f32"; - triangular.recurrenceStartRow = 2; - triangular.recurrenceLoopCount = 4; - triangular.denseDotTailOps = 16; - triangular.requiresCubeTailPartition = true; - f.simtAnchors.triangularSolves.push_back(triangular); - return f; -} - } // namespace namespace { @@ -331,14 +100,12 @@ logicalStage(llvm::StringRef id, StageCostModelKind kind, int64_t iterations = 1) { LogicalStage stage; stage.id = id.str(); - stage.description = id.str(); stage.costModelKind = kind; stage.scheduleKind = schedule; stage.iterationCount = iterations; stage.simdLegal = true; stage.simtLegal = true; stage.legalSimtFactors = {1}; - stage.features.source = "unit_test_stage_analysis"; stage.workload.paysKernelSetup = true; stage.workload.operationElements["f32.add"] = 64.0; stage.workload.issueElements = 4.0; @@ -350,7 +117,6 @@ evaluateOneStage(LogicalStage stage, HardwareProfile profile = hardwareProfile()) { LogicalPhase phase; phase.id = "phase"; - phase.description = "phase"; phase.stages.push_back(std::move(stage)); StagePartition partition; partition.domain = "unit_test"; @@ -360,122 +126,6 @@ evaluateOneStage(LogicalStage stage, } // namespace -TEST(SimdSimtCostModelTest, GatherDotUsesStageCostEvaluator) { - auto report = estimateSimdSimtCandidates(gatherDotFeatures(), options(32)); - if (!report) - FAIL() << llvm::toString(report.takeError()); - - EXPECT_TRUE(report->stageModel.applied); - EXPECT_EQ(report->stageModel.domain, "indirect_underfilled_dot"); - EXPECT_GT(report->candidateCosts.allSimd, 0.0); - EXPECT_GT(report->candidateCosts.allSimtOnly, 0.0); - EXPECT_GT(report->candidateCosts.mixedSimdSimt, 0.0); - for (const auto &stage : report->stageModel.stages) - for (const auto &cost : stage.implementations) { - EXPECT_TRUE(cost.implementation.mode == StageMode::SIMD || - cost.implementation.mode == StageMode::SIMT); - EXPECT_GT(cost.totalCycles, 0.0) << stage.id; - EXPECT_FALSE(cost.modelName.empty()); - EXPECT_EQ(cost.profileVersion, report->profileVersion); - } -} - -TEST(SimdSimtCostModelTest, FbgemmWithoutAutoBlockifyExposesOnlySuperBlockF1) { - auto report = estimateSimdSimtCandidates(fbgemmFeatures(), options(4)); - if (!report) - FAIL() << llvm::toString(report.takeError()); - - EXPECT_TRUE(report->stageModel.applied); - EXPECT_EQ(report->stageModel.domain, "loaded_index_rowwise_reduction"); - EXPECT_TRUE(report->stageModel.allSimd.legal); - EXPECT_TRUE(report->stageModel.allSimt.legal); - EXPECT_TRUE(report->stageModel.mixed.legal); - EXPECT_EQ(report->breakdown.mixedCostSource, - "stage_cost_evaluator_route_sum"); - auto reduction = - llvm::find_if(report->stageModel.stages, [](const auto &stage) { - return stage.model == "rowwise_reduction"; - }); - ASSERT_NE(reduction, report->stageModel.stages.end()); - ASSERT_EQ(reduction->implementations.size(), 2u); - EXPECT_EQ(reduction->implementations[0].modelName, "simd_reduction"); - EXPECT_EQ(reduction->implementations[1].modelName, "simt_reduction"); - EXPECT_EQ(reduction->implementations[1].implementation.superblockFactor, 1); - EXPECT_GT(reduction->implementations[0].resources.criticalPath, 0.0); -} - -TEST(SimdSimtCostModelTest, TriangularSolveUsesStageCostEvaluator) { - auto report = - estimateSimdSimtCandidates(triangularUnknownLoopFeatures(), options(32)); - if (!report) - FAIL() << llvm::toString(report.takeError()); - - EXPECT_TRUE(report->stageModel.applied); - EXPECT_EQ(report->stageModel.domain, "triangular_recurrence"); - ASSERT_EQ(report->stageModel.phases.size(), 4u); - ASSERT_EQ(report->stageModel.stages.size(), 5u); - ASSERT_EQ(report->stageModel.mixed.implementations.size(), 5u); - bool hasSimd = false; - bool hasSimt = false; - for (const auto &implementation : report->stageModel.mixed.implementations) { - hasSimd |= implementation.mode == StageMode::SIMD; - hasSimt |= implementation.mode == StageMode::SIMT; - } - EXPECT_TRUE(hasSimd); - EXPECT_TRUE(hasSimt); - EXPECT_EQ(report->stageModel.mixed.routeSuperblockFactor, 1); - EXPECT_EQ(report->breakdown.mixedCostSource, - "stage_cost_evaluator_route_sum"); -} - -TEST(SimdSimtCostModelTest, Bt16RecurrenceEvaluatesF2AndF4Pressure) { - auto features = triangularBt16StageFeatures(); - features.autoBlockifyV1Applied = true; - auto report = estimateSimdSimtCandidates(features, options(4)); - if (!report) - FAIL() << llvm::toString(report.takeError()); - ASSERT_TRUE(report->stageModel.applied); - EXPECT_EQ(report->stageModel.domain, "triangular_recurrence"); - auto recurrence = - llvm::find_if(report->stageModel.stages, [](const auto &stage) { - return stage.model == "loop_carried_recurrence"; - }); - ASSERT_NE(recurrence, report->stageModel.stages.end()); - ASSERT_EQ(recurrence->iterationCount, 14); - auto findFactor = [&](int64_t factor) { - return llvm::find_if( - recurrence->implementations, [&](const auto &implementation) { - return implementation.implementation.mode == StageMode::SIMT && - implementation.implementation.superblockFactor == factor; - }); - }; - const auto f2 = findFactor(2); - const auto f4 = findFactor(4); - ASSERT_NE(f2, recurrence->implementations.end()); - ASSERT_NE(f4, recurrence->implementations.end()); - // Feature-summary fallback has no exact SSA live-out byte count, so it can - // model F4 latency hiding but cannot invent recurrence-state pressure. The - // operation-graph test below supplies liveOutBytes and verifies F4 > F2. - EXPECT_LE(f4->totalCycles, f2->totalCycles); -} - -TEST(SimdSimtCostModelTest, UnknownPatternUsesAggregateFallback) { - auto features = triangularUnknownLoopFeatures(); - features.simtAnchors.mechanismKinds.clear(); - features.simtAnchors.triangularSolves.clear(); - features.simtAnchors.count = 0; - auto report = estimateSimdSimtCandidates(features, options(32)); - if (!report) - FAIL() << llvm::toString(report.takeError()); - EXPECT_FALSE(report->stageModel.applied); - EXPECT_GT(report->candidateCosts.allSimd, 0.0); -} - -TEST(SimdSimtCostModelTest, RegistryCoversEveryModeAndKind) { - if (llvm::Error error = StageCostModelRegistry::get().verifyComplete()) - FAIL() << llvm::toString(std::move(error)); -} - TEST(SimdSimtCostModelTest, StageHasOnlySimdOrSimtImplementations) { LogicalStage stage = logicalStage("scalar", StageCostModelKind::ScalarIssue); auto table = evaluateOneStage(std::move(stage)); @@ -489,14 +139,18 @@ TEST(SimdSimtCostModelTest, StageHasOnlySimdOrSimtImplementations) { } TEST(SimdSimtCostModelTest, - ScopeSuperBlockLegalityOpensOnlyWithBackendMaterializer) { - auto makePartition = [] { + ScopeSuperBlockLegalityRequiresBackendAndResourceMaximum) { + auto makePartition = [](int64_t independentGroups) { StagePartition partition; partition.domain = "unit_test"; mlir::ascend::LogicalPhase phase; phase.id = "phase"; LogicalStage stage = - logicalStage("payload", StageCostModelKind::ScalarIssue); + logicalStage("payload", StageCostModelKind::LoopCarriedRecurrence, + StageScheduleKind::LoopCarriedSerial, /*iterations=*/16); + stage.features.hasLoop = true; + stage.features.hasLoopCarriedDataDependency = true; + stage.features.parallelRecurrenceGroupCount = independentGroups; stage.localSimtMaterializable = true; stage.localSimtFactors = {1}; phase.stages.push_back(std::move(stage)); @@ -504,18 +158,67 @@ TEST(SimdSimtCostModelTest, return partition; }; - StagePartition f1Only = makePartition(); + StagePartition f1Only = makePartition(/*independentGroups=*/4); if (llvm::Error error = StageModeLegalityAnalysis().analyze(f1Only, 4, false)) FAIL() << llvm::toString(std::move(error)); EXPECT_EQ(f1Only.phases[0].stages[0].localSimtFactors, (std::vector{1})); - StagePartition scopeSuperblock = makePartition(); + StagePartition scopeSuperblock = makePartition(/*independentGroups=*/4); if (llvm::Error error = StageModeLegalityAnalysis().analyze(scopeSuperblock, 4, true)) FAIL() << llvm::toString(std::move(error)); EXPECT_EQ(scopeSuperblock.phases[0].stages[0].localSimtFactors, (std::vector{1, 2, 4})); + + // ABI-v2 creates an F1 V1 scheduling loop and refines only the selected + // scope after bufferization, but local and whole-kernel factors still share + // the same target/runtime warp-resource maximum. + StagePartition mixedOnly = makePartition(/*independentGroups=*/4); + if (llvm::Error error = + StageModeLegalityAnalysis().analyze(mixedOnly, 1, true)) + FAIL() << llvm::toString(std::move(error)); + EXPECT_EQ(mixedOnly.phases[0].stages[0].legalSimtFactors, + (std::vector{1})); + EXPECT_EQ(mixedOnly.phases[0].stages[0].localSimtFactors, + (std::vector{1})); + auto mixedOnlyCosts = evaluateOneStage(mixedOnly.phases[0].stages[0]); + if (!mixedOnlyCosts) + FAIL() << llvm::toString(mixedOnlyCosts.takeError()); + ASSERT_EQ(mixedOnlyCosts->stages[0].implementations.size(), 3u); + EXPECT_EQ(mixedOnlyCosts->stages[0].legalSimtFactors, + (std::vector{1})); + EXPECT_EQ(mixedOnlyCosts->stages[0].localSimtFactors, + (std::vector{1})); + + StagePartition oneWorkGroup = makePartition(/*independentGroups=*/1); + if (llvm::Error error = + StageModeLegalityAnalysis().analyze(oneWorkGroup, 4, true)) + FAIL() << llvm::toString(std::move(error)); + EXPECT_EQ(oneWorkGroup.phases[0].stages[0].localSimtFactors, + (std::vector{1, 2, 4})); +} + +TEST(SimdSimtCostModelTest, LocalScopeFactorsHonorKernelResourceMaximum) { + StagePartition partition; + partition.domain = "local_factor_limit"; + LogicalPhase phase; + phase.id = "gather"; + LogicalStage stage; + stage.id = "indirect_tile_gather"; + stage.costModelKind = StageCostModelKind::IndirectGatherMemory; + stage.scheduleKind = StageScheduleKind::PartiallyDependent; + stage.iterationCount = 1; + stage.localSimtMaterializable = true; + phase.stages.push_back(std::move(stage)); + partition.phases.push_back(std::move(phase)); + + ASSERT_FALSE(StageModeLegalityAnalysis().analyze( + partition, /*maximumSuperblockFactor=*/2, + /*scopeSuperblockMaterializable=*/true)); + const LogicalStage &result = partition.phases.front().stages.front(); + EXPECT_EQ(result.legalSimtFactors, (std::vector{1, 2})); + EXPECT_EQ(result.localSimtFactors, (std::vector{1, 2})); } TEST(SimdSimtCostModelTest, KernelMixedRouteComesFromAdjacentStageModes) { @@ -525,20 +228,17 @@ TEST(SimdSimtCostModelTest, KernelMixedRouteComesFromAdjacentStageModes) { auto addStage = [&](llvm::StringRef id, double simd, double simt) { mlir::ascend::LogicalStageCost stage; stage.id = id.str(); - stage.description = id.str(); stage.localSimtMaterializable = true; stage.localSimtFactors = {1}; - auto cost = [&](StageMode mode, double cycles) { + auto cost = [&](StageMode mode, double cycles, bool localScope = false) { mlir::ascend::StageImplementationCost result; - result.implementation = {mode, 1}; + result.implementation = {mode, 1, localScope}; result.totalCycles = cycles; - result.modelName = "unit_test"; - result.profileVersion = table.profileVersion; - result.source = "unit_test"; return result; }; stage.implementations = {cost(StageMode::SIMD, simd), - cost(StageMode::SIMT, simt)}; + cost(StageMode::SIMT, simt), + cost(StageMode::SIMT, simt, true)}; table.stages.push_back(stage); }; addStage("head", 10.0, 20.0); @@ -562,39 +262,36 @@ TEST(SimdSimtCostModelTest, KernelMixedRouteComesFromAdjacentStageModes) { EXPECT_EQ(result->mixed.implementations[0].mode, StageMode::SIMD); EXPECT_EQ(result->mixed.implementations[1].mode, StageMode::SIMT); EXPECT_EQ(result->mixed.implementations[2].mode, StageMode::SIMD); + ASSERT_EQ(result->mixed.entryTransitionCycles.size(), 3u); + EXPECT_DOUBLE_EQ(result->mixed.entryTransitionCycles[1], 12.0); } TEST(SimdSimtCostModelTest, MixedScopePaysExactBidirectionalUbHandoffCost) { StageCostTable table; table.domain = "scope_handoff"; table.profileVersion = "unit-test-profile-v1"; - auto makeCost = [&](StageMode mode, double cycles) { + auto makeCost = [&](StageMode mode, double cycles, bool localScope = false) { mlir::ascend::StageImplementationCost cost; - cost.implementation = {mode, 1}; + cost.implementation = {mode, 1, localScope}; cost.totalCycles = cycles; - cost.modelName = "unit_test"; - cost.profileVersion = table.profileVersion; - cost.source = "unit_test"; return cost; }; mlir::ascend::LogicalStageCost head; head.id = "head"; - head.description = head.id; head.implementations = {makeCost(StageMode::SIMD, 10.0), makeCost(StageMode::SIMT, 20.0)}; mlir::ascend::LogicalStageCost payload; payload.id = "large_result_payload"; - payload.description = payload.id; payload.localSimtMaterializable = true; payload.localSimtFactors = {1}; payload.localSimtScopeCount = 2; payload.scopeInputTensorBytes = 4096; payload.scopeOutputTensorBytes = 16384; payload.implementations = {makeCost(StageMode::SIMD, 100.0), - makeCost(StageMode::SIMT, 10.0)}; + makeCost(StageMode::SIMT, 10.0), + makeCost(StageMode::SIMT, 10.0, true)}; mlir::ascend::LogicalStageCost tail = head; tail.id = "tail"; - tail.description = tail.id; table.stages = {head, payload, tail}; mlir::ascend::LogicalPhaseCost phase; phase.id = "phase"; @@ -615,9 +312,72 @@ TEST(SimdSimtCostModelTest, MixedScopePaysExactBidirectionalUbHandoffCost) { // Output: 16384/(4*32) + 16384/512 = 160 cycles. // Head/payload/tail: 10 + (10 + 208) + 10 = 238 cycles. EXPECT_DOUBLE_EQ(routes->mixed.totalCycles, 238.0); + ASSERT_EQ(routes->mixed.entryTransitionCycles.size(), 3u); + EXPECT_DOUBLE_EQ(routes->mixed.entryTransitionCycles[0], 0.0); + EXPECT_DOUBLE_EQ(routes->mixed.entryTransitionCycles[1], 208.0); + EXPECT_DOUBLE_EQ(routes->mixed.entryTransitionCycles[2], 0.0); EXPECT_GT(routes->mixed.totalCycles, routes->allSimd.totalCycles); } +TEST(SimdSimtCostModelTest, MixedScopeSuperBlockAmortizesOnlyFixedTransitions) { + StageCostTable table; + table.domain = "scope_superblock_handoff"; + table.profileVersion = "unit-test-profile-v1"; + auto makeCost = [&](StageMode mode, int64_t factor, double cycles, + bool localScope = false) { + StageImplementationCost cost; + cost.implementation = {mode, factor, localScope}; + cost.totalCycles = cycles; + return cost; + }; + + LogicalStageCost head; + head.id = "head"; + head.implementations = {makeCost(StageMode::SIMD, 1, 10.0), + makeCost(StageMode::SIMT, 1, 100.0)}; + + LogicalStageCost payload; + payload.id = "payload"; + payload.localSimtMaterializable = true; + payload.localSimtFactors = {1, 2, 4}; + payload.localSimtScopeCount = 1; + payload.scopeInputTensorBytes = 4096; + payload.scopeOutputTensorBytes = 4096; + payload.implementations = {makeCost(StageMode::SIMD, 1, 1000.0), + makeCost(StageMode::SIMT, 1, 100.0), + makeCost(StageMode::SIMT, 1, 100.0, true), + makeCost(StageMode::SIMT, 2, 100.0, true), + makeCost(StageMode::SIMT, 4, 100.0, true)}; + + LogicalStageCost tail = head; + tail.id = "tail"; + table.stages = {head, payload, tail}; + LogicalPhaseCost phase; + phase.id = "phase"; + phase.stages = table.stages; + table.phases.push_back(std::move(phase)); + + StageTransitionCost transition; + transition.simdToSimtCycles = 40.0; + transition.simtToSimdCycles = 40.0; + transition.simdUbLoadBytesPerCycle = 512.0; + transition.simdUbStoreBytesPerCycle = 256.0; + transition.simtUbLoadBytesPerThreadPerCycle = 4.0; + transition.simtUbStoreBytesPerThreadPerCycle = 4.0; + transition.simtWarpSize = 32; + auto routes = solveStageRoutes(table, transition); + if (!routes) + FAIL() << llvm::toString(routes.takeError()); + + ASSERT_TRUE(routes->mixed.legal); + EXPECT_EQ(routes->mixed.routeSuperblockFactor, 4); + // The 80-cycle fixed transition pair is amortized to 20 cycles. The + // 4096-byte input/output handoff remains 48 + 40 cycles per program. + EXPECT_DOUBLE_EQ(routes->mixed.totalCycles, + 10.0 + 100.0 + 20.0 + 48.0 + 40.0 + 10.0); + EXPECT_DOUBLE_EQ(routes->mixed.entryTransitionCycles[1], 108.0); +} + TEST(SimdSimtCostModelTest, IndependentLoopUsesSimdRooflineAndSerialSimtCost) { LogicalStage stage = logicalStage("independent", StageCostModelKind::IndependentPipelinedLoop, @@ -638,6 +398,34 @@ TEST(SimdSimtCostModelTest, IndependentLoopUsesSimdRooflineAndSerialSimtCost) { table->stages[0].implementations[1].totalCycles); } +TEST(SimdSimtCostModelTest, SimdCubeStageUsesCvPipelineCriticalPath) { + LogicalStage stage = + logicalStage("cube", StageCostModelKind::TinyCubeRoofline, + StageScheduleKind::IndependentPipelined, 4); + stage.features.hasDot = true; + stage.features.hasContiguousMemory = true; + stage.workload.operationElements.clear(); + stage.workload.scalarOperations = 16.0; + stage.workload.loadBytes = 2048.0; + stage.workload.storeBytes = 512.0; + stage.workload.dotFlops = 8192.0; + stage.workload.issueElements = 128.0; + + auto table = evaluateOneStage(stage); + if (!table) + FAIL() << llvm::toString(table.takeError()); + + const StageImplementationCost &simd = table->stages[0].implementations[0]; + ASSERT_EQ(simd.implementation.mode, StageMode::SIMD); + EXPECT_DOUBLE_EQ(simd.resources.setup, 18.0); + EXPECT_DOUBLE_EQ(simd.resources.load, 64.0); + EXPECT_DOUBLE_EQ(simd.resources.store, 32.0); + EXPECT_DOUBLE_EQ(simd.resources.dot, 128.0); + // SIMD/Cube CV resources overlap inside one Stage. The model charges the + // critical resource (MMAD here), rather than serializing load+MMAD+store. + EXPECT_DOUBLE_EQ(simd.totalCycles, 18.0 + 4.0 * (16.0 + 128.0)); +} + TEST(SimdSimtCostModelTest, TrueLoopCarriedDependencyDisablesSimdRoofline) { LogicalStage stage = logicalStage("dependent", StageCostModelKind::IndependentPipelinedLoop, @@ -731,12 +519,11 @@ TEST(SimdSimtCostModelTest, groupedTable->stages[0].implementations[0].totalCycles; EXPECT_LT(groupedCycles, serialCycles); const auto &resources = groupedTable->stages[0].implementations[0].resources; - EXPECT_GE(groupedCycles, - resources.setup + 16.0 * resources.issue + resources.epilogue); + EXPECT_GE(groupedCycles, resources.setup + 16.0 * resources.issue); } TEST(SimdSimtCostModelTest, - SuperBlockF4ChargesReplicatedPersistentRecurrenceState) { + SuperBlockRecurrenceContentionStartsAbovePressureFreeFactor) { LogicalStage stage = logicalStage("stateful_recurrence", StageCostModelKind::LoopCarriedRecurrence, StageScheduleKind::LoopCarriedSerial, 16); @@ -748,6 +535,7 @@ TEST(SimdSimtCostModelTest, stage.liveOutBytes = 4096; stage.workload.loadWarpInstructions = 16.0; stage.workload.shuffleLaneSteps = 128.0; + stage.workload.issueElements = 64.0; HardwareProfile profile = hardwareProfile(); profile.logicalWarpGroupCount = 4; @@ -756,8 +544,48 @@ TEST(SimdSimtCostModelTest, FAIL() << llvm::toString(table.takeError()); const auto &costs = table->stages.front().implementations; ASSERT_EQ(costs.size(), 3u); - EXPECT_LT(costs[1].totalCycles, costs[0].totalCycles); - EXPECT_GT(costs[2].totalCycles, costs[1].totalCycles); + EXPECT_LE(costs[1].totalCycles, costs[0].totalCycles); + EXPECT_LE(costs[2].totalCycles, costs[1].totalCycles); + EXPECT_GE(costs[2].totalCycles, + costs[2].resources.setup + 16.0 * costs[2].resources.issue); + + stage.workload.estimatedSpillTransactions = 32.0; + auto spillingTable = evaluateOneStage(stage, profile); + if (!spillingTable) + FAIL() << llvm::toString(spillingTable.takeError()); + const auto &spillingCosts = spillingTable->stages.front().implementations; + ASSERT_EQ(spillingCosts.size(), 3u); + EXPECT_GT(spillingCosts[2].totalCycles, spillingCosts[1].totalCycles); +} + +TEST(SimdSimtCostModelTest, + MixedLocalStageSuperBlockHidesRecurrenceLatencyButKeepsIssueFloor) { + LogicalStage stage = logicalStage( + "mixed_recurrence", StageCostModelKind::LoopCarriedRecurrence, + StageScheduleKind::LoopCarriedSerial, /*iterations=*/16); + stage.simdLegal = false; + stage.legalSimtFactors = {1}; + stage.localSimtMaterializable = true; + stage.localSimtFactors = {1, 2, 4}; + stage.features.hasLoop = true; + stage.features.hasLoopCarriedDataDependency = true; + stage.workload.loadWarpInstructions = 16.0; + stage.workload.shuffleLaneSteps = 128.0; + stage.workload.issueElements = 64.0; + + auto table = evaluateOneStage(stage, hardwareProfile()); + if (!table) + FAIL() << llvm::toString(table.takeError()); + const auto &costs = table->stages.front().implementations; + ASSERT_EQ(costs.size(), 4u); + EXPECT_FALSE(costs[0].implementation.localScope); + EXPECT_TRUE(costs[1].implementation.localScope); + EXPECT_TRUE(costs[2].implementation.localScope); + EXPECT_TRUE(costs[3].implementation.localScope); + EXPECT_LE(costs[2].totalCycles, costs[1].totalCycles); + EXPECT_LE(costs[3].totalCycles, costs[2].totalCycles); + EXPECT_GE(costs[3].totalCycles, + costs[3].resources.setup + 16.0 * costs[3].resources.issue); } TEST(SimdSimtCostModelTest, IndirectMemoryUsesDependencyProfile) { @@ -787,24 +615,19 @@ TEST(SimdSimtCostModelTest, MixedRouteRejectsUnmaterializableSimtStage) { StageCostTable table; table.domain = "unit_test"; table.profileVersion = "unit-test-profile-v1"; - auto makeCost = [&](StageMode mode, double cycles) { + auto makeCost = [&](StageMode mode, double cycles, bool localScope = false) { mlir::ascend::StageImplementationCost cost; - cost.implementation = {mode, 1}; + cost.implementation = {mode, 1, localScope}; cost.totalCycles = cycles; - cost.modelName = "unit_test"; - cost.profileVersion = table.profileVersion; - cost.source = "unit_test"; return cost; }; mlir::ascend::LogicalStageCost head; head.id = "head"; - head.description = head.id; head.localSimtMaterializable = false; head.implementations = {makeCost(StageMode::SIMD, 1.0), makeCost(StageMode::SIMT, 100.0)}; mlir::ascend::LogicalStageCost payload; payload.id = "unmaterializable_payload"; - payload.description = payload.id; payload.localSimtMaterializable = false; payload.implementations = {makeCost(StageMode::SIMD, 100.0), makeCost(StageMode::SIMT, 1.0)}; @@ -820,35 +643,112 @@ TEST(SimdSimtCostModelTest, MixedRouteRejectsUnmaterializableSimtStage) { EXPECT_FALSE(routes->mixed.legal); } +TEST(SimdSimtCostModelTest, + MixedRouteReportsCheapestConstrainedRouteWhenLocalScopeLoses) { + StageCostTable table; + table.domain = "constrained_mixed"; + table.profileVersion = "unit-test-profile-v1"; + auto makeCost = [&](StageMode mode, double cycles, bool localScope = false) { + mlir::ascend::StageImplementationCost cost; + cost.implementation = {mode, 1, localScope}; + cost.totalCycles = cycles; + return cost; + }; + + mlir::ascend::LogicalStageCost gather; + gather.id = "indirect_tile_gather"; + gather.localSimtMaterializable = true; + gather.localSimtScopeCount = 1; + gather.implementations = {makeCost(StageMode::SIMD, 100.0), + makeCost(StageMode::SIMT, 130.0), + makeCost(StageMode::SIMT, 130.0, true)}; + mlir::ascend::LogicalStageCost dot; + dot.id = "tiny_cube_dot"; + dot.implementations = {makeCost(StageMode::SIMD, 40.0), + makeCost(StageMode::SIMT, 90.0)}; + table.stages = {gather, dot}; + mlir::ascend::LogicalPhaseCost phase; + phase.id = "gather_dot_min"; + phase.stages = table.stages; + table.phases.push_back(std::move(phase)); + + StageTransitionCost transition; + transition.simdToSimtCycles = 10.0; + transition.simtToSimdCycles = 10.0; + auto routes = solveStageRoutes(table, transition); + if (!routes) + FAIL() << llvm::toString(routes.takeError()); + ASSERT_TRUE(routes->mixed.legal); + ASSERT_EQ(routes->mixed.implementations.size(), 2u); + EXPECT_EQ(routes->mixed.implementations[0].mode, StageMode::SIMT); + EXPECT_TRUE(routes->mixed.implementations[0].localScope); + EXPECT_EQ(routes->mixed.implementations[1].mode, StageMode::SIMD); + EXPECT_DOUBLE_EQ(routes->mixed.totalCycles, 190.0); +} + +TEST(SimdSimtCostModelTest, AllSimdDoesNotPayRouteConditionalAutoBlockify) { + StageCostTable table; + table.domain = "auto_blockify_route_conditional"; + table.profileVersion = "unit-test-profile-v1"; + auto makeCost = [&](StageMode mode, double cycles) { + mlir::ascend::StageImplementationCost cost; + cost.implementation = {mode, 1, false}; + cost.totalCycles = cycles; + return cost; + }; + + mlir::ascend::LogicalStageCost dispatch; + dispatch.id = "physical_program_dispatch"; + dispatch.model = "auto_blockify_dispatch"; + dispatch.implementations = {makeCost(StageMode::SIMD, 40.0), + makeCost(StageMode::SIMT, 30.0)}; + mlir::ascend::LogicalStageCost payload; + payload.id = "payload"; + payload.model = "scalar_issue"; + payload.implementations = {makeCost(StageMode::SIMD, 100.0), + makeCost(StageMode::SIMT, 80.0)}; + table.stages = {dispatch, payload}; + mlir::ascend::LogicalPhaseCost phase; + phase.id = "phase"; + phase.stages = table.stages; + table.phases.push_back(std::move(phase)); + + auto routes = solveStageRoutes(table, StageTransitionCost{}); + if (!routes) + FAIL() << llvm::toString(routes.takeError()); + ASSERT_TRUE(routes->allSimd.legal); + ASSERT_EQ(routes->allSimd.logicalStageCycles.size(), 2u); + EXPECT_DOUBLE_EQ(routes->allSimd.logicalStageCycles[0], 0.0); + EXPECT_DOUBLE_EQ(routes->allSimd.logicalStageCycles[1], 100.0); + EXPECT_DOUBLE_EQ(routes->allSimd.totalCycles, 100.0); + ASSERT_TRUE(routes->allSimt.legal); + EXPECT_DOUBLE_EQ(routes->allSimt.totalCycles, 110.0); +} + TEST(SimdSimtCostModelTest, MixedRouteChargesEveryMaterializedScope) { StageCostTable table; table.domain = "scope_count"; table.profileVersion = "unit-test-profile-v1"; - auto makeCost = [&](StageMode mode, double cycles) { + auto makeCost = [&](StageMode mode, double cycles, bool localScope = false) { mlir::ascend::StageImplementationCost cost; - cost.implementation = {mode, 1}; + cost.implementation = {mode, 1, localScope}; cost.totalCycles = cycles; - cost.modelName = "unit_test"; - cost.profileVersion = table.profileVersion; - cost.source = "unit_test"; return cost; }; mlir::ascend::LogicalStageCost head; head.id = "head"; - head.description = head.id; head.implementations = {makeCost(StageMode::SIMD, 1.0), makeCost(StageMode::SIMT, 100.0)}; mlir::ascend::LogicalStageCost gather; gather.id = "two_anchor_gather"; - gather.description = gather.id; gather.localSimtMaterializable = true; gather.localSimtFactors = {1}; gather.localSimtScopeCount = 2; gather.implementations = {makeCost(StageMode::SIMD, 100.0), - makeCost(StageMode::SIMT, 1.0)}; + makeCost(StageMode::SIMT, 1.0), + makeCost(StageMode::SIMT, 1.0, true)}; mlir::ascend::LogicalStageCost tail = head; tail.id = "tail"; - tail.description = tail.id; table.stages = {head, gather, tail}; mlir::ascend::LogicalPhaseCost phase; phase.id = "phase"; @@ -865,6 +765,8 @@ TEST(SimdSimtCostModelTest, MixedRouteChargesEveryMaterializedScope) { // 1 SIMD head + (10 enter + 1 payload + 20 extra scope pair) + // (10 leave + 1 SIMD tail). EXPECT_DOUBLE_EQ(routes->mixed.totalCycles, 43.0); + ASSERT_EQ(routes->mixed.entryTransitionCycles.size(), 3u); + EXPECT_DOUBLE_EQ(routes->mixed.entryTransitionCycles[1], 40.0); } TEST(SimdSimtCostModelTest, SuperBlockLatencyHidingStopsAtUsefulFactorLimit) { @@ -879,13 +781,15 @@ TEST(SimdSimtCostModelTest, SuperBlockLatencyHidingStopsAtUsefulFactorLimit) { auto table = evaluateOneStage(stage, cappedProfile); if (!table) FAIL() << llvm::toString(table.takeError()); + table->logicalProgramCountHint = 64; + table->physicalCoreCountHint = 32; auto routes = solveStageRoutes(*table, cappedProfile.transition); if (!routes) FAIL() << llvm::toString(routes.takeError()); EXPECT_TRUE(routes->allSimt.legal); EXPECT_EQ(routes->allSimt.routeSuperblockFactor, 2); EXPECT_LT(routes->allSimt.totalCycles, - table->stages[0].implementations[0].totalCycles); + 2.0 * table->stages[0].implementations[0].totalCycles); } TEST(SimdSimtCostModelTest, PureSimtRouteUsesOneUniformSuperBlockFactor) { @@ -896,19 +800,14 @@ TEST(SimdSimtCostModelTest, PureSimtRouteUsesOneUniformSuperBlockFactor) { mlir::ascend::StageImplementationCost cost; cost.implementation = {StageMode::SIMT, factor}; cost.totalCycles = cycles; - cost.modelName = "unit_test"; - cost.profileVersion = table.profileVersion; - cost.source = "unit_test"; return cost; }; mlir::ascend::LogicalStageCost first; first.id = "first"; - first.description = first.id; first.implementations = {makeCost(1, 5.0), makeCost(2, 1.0), makeCost(4, 3.0)}; mlir::ascend::LogicalStageCost second; second.id = "second"; - second.description = second.id; second.implementations = {makeCost(1, 5.0), makeCost(2, 4.0), makeCost(4, 1.0)}; table.stages = {first, second}; @@ -932,30 +831,34 @@ TEST(SimdSimtCostModelTest, MixedScopeSuperBlockUsesSelectedFactorCost) { StageCostTable table; table.domain = "mixed_scope_superblock"; table.profileVersion = "unit-test-profile-v1"; - auto makeCost = [&](StageMode mode, int64_t factor, double cycles) { + auto makeCost = [&](StageMode mode, int64_t factor, double cycles, + bool localScope = false) { mlir::ascend::StageImplementationCost cost; - cost.implementation = {mode, factor}; + cost.implementation = {mode, factor, localScope}; cost.totalCycles = cycles; - cost.modelName = "unit_test"; - cost.profileVersion = table.profileVersion; - cost.source = "unit_test"; return cost; }; mlir::ascend::LogicalStageCost prefix; prefix.id = "simd_prefix"; - prefix.description = prefix.id; - prefix.implementations = { - makeCost(StageMode::SIMD, 1, 5.0), makeCost(StageMode::SIMT, 1, 50.0), - makeCost(StageMode::SIMT, 2, 25.0), makeCost(StageMode::SIMT, 4, 12.5)}; + prefix.implementations = {makeCost(StageMode::SIMD, 1, 5.0), + makeCost(StageMode::SIMT, 1, 50.0), + makeCost(StageMode::SIMT, 2, 25.0), + makeCost(StageMode::SIMT, 4, 12.5), + makeCost(StageMode::SIMT, 1, 50.0, true), + makeCost(StageMode::SIMT, 2, 25.0, true), + makeCost(StageMode::SIMT, 4, 12.5, true)}; prefix.localSimtMaterializable = true; prefix.localSimtFactors = {1, 2, 4}; mlir::ascend::LogicalStageCost payload; payload.id = "local_simt_payload"; - payload.description = payload.id; - payload.implementations = { - makeCost(StageMode::SIMD, 1, 100.0), makeCost(StageMode::SIMT, 1, 10.0), - makeCost(StageMode::SIMT, 2, 1.0), makeCost(StageMode::SIMT, 4, 0.5)}; + payload.implementations = {makeCost(StageMode::SIMD, 1, 100.0), + makeCost(StageMode::SIMT, 1, 10.0), + makeCost(StageMode::SIMT, 2, 1.0), + makeCost(StageMode::SIMT, 4, 0.5), + makeCost(StageMode::SIMT, 1, 10.0, true), + makeCost(StageMode::SIMT, 2, 1.0, true), + makeCost(StageMode::SIMT, 4, 0.5, true)}; payload.localSimtMaterializable = true; payload.localSimtFactors = {1, 2, 4}; table.stages = {prefix, payload}; @@ -972,31 +875,6 @@ TEST(SimdSimtCostModelTest, MixedScopeSuperBlockUsesSelectedFactorCost) { EXPECT_DOUBLE_EQ(routes->mixed.totalCycles, 5.5); } -TEST(SimdSimtCostModelTest, StageWorkloadAnalysisConservesKernelWork) { - const SimdSimtFeatureSummary features = gatherDotFeatures(); - auto result = - StagePartitioner().partition(features, StagePartitionerOptions{}); - if (!result) - FAIL() << llvm::toString(result.takeError()); - ASSERT_TRUE(*result); - const StagePartition &partition = **result; - double operations = 0.0, bytes = 0.0, warps = 0.0, flops = 0.0, issue = 0.0; - for (const auto &phase : partition.phases) - for (const auto &stage : phase.stages) { - for (const auto &entry : stage.workload.operationElements) - operations += entry.second * stage.iterationCount; - bytes += stage.workload.loadBytes * stage.iterationCount; - warps += stage.workload.loadWarpInstructions * stage.iterationCount; - flops += stage.workload.dotFlops * stage.iterationCount; - issue += stage.workload.issueElements * stage.iterationCount; - } - EXPECT_DOUBLE_EQ(operations, 82.0); - EXPECT_DOUBLE_EQ(bytes, 1088.0); - EXPECT_DOUBLE_EQ(warps, 17.0); - EXPECT_DOUBLE_EQ(flops, 8192.0); - EXPECT_GT(issue, 0.0); -} - TEST(SimdSimtCostModelTest, OperationGraphBoundaryOwnsEveryRootAndDerivesLiveValues) { mlir::MLIRContext context; @@ -1033,9 +911,9 @@ TEST(SimdSimtCostModelTest, anchor.scopeOperations.push_back(recurrence); anchor.scopeInsertionPoint = recurrence; anchor.kind = mlir::ascend::SimtAnchorKind::TriangularSolveLoop; - anchor.facts = + anchor.triangularSolve = triangularBt16StageFeatures().simtAnchors.triangularSolves.front(); - anchor.lowerability.mixed = mlir::ascend::CandidateLoweringStatus::Native; + anchor.lowerability.mixed = true; anchor.materializable = true; mlir::ascend::SimtAnchorPlan anchorPlan; anchorPlan.anchors.push_back(std::move(anchor)); @@ -1058,7 +936,6 @@ TEST(SimdSimtCostModelTest, FAIL() << llvm::toString(result.takeError()); ASSERT_TRUE(*result); const StagePartition &partition = **result; - EXPECT_EQ(partition.boundarySource, "operation_graph"); EXPECT_TRUE(partition.operationOwnershipComplete); int64_t ownedRootCount = 0; @@ -1131,9 +1008,9 @@ TEST(SimdSimtCostModelTest, anchor.scopeOperations = {setup, recurrence}; anchor.scopeInsertionPoint = recurrence; anchor.kind = mlir::ascend::SimtAnchorKind::TriangularSolveLoop; - anchor.facts = + anchor.triangularSolve = triangularBt16StageFeatures().simtAnchors.triangularSolves.front(); - anchor.lowerability.mixed = mlir::ascend::CandidateLoweringStatus::Native; + anchor.lowerability.mixed = true; anchor.materializable = true; mlir::ascend::SimtAnchorPlan anchorPlan; anchorPlan.anchors.push_back(std::move(anchor)); @@ -1171,6 +1048,69 @@ TEST(SimdSimtCostModelTest, EXPECT_EQ(recurrenceStage->simtAnchorIndices, std::vector({0})); } +TEST(SimdSimtCostModelTest, + LocalScopeReturningPointerTensorIsRejectedBeforeScoring) { + mlir::MLIRContext context; + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.allowUnregisteredDialects(); + auto module = mlir::parseSourceString(R"mlir( + module { + func.func @kernel(%base: !tt.ptr) { + %indices = arith.constant dense<0> : tensor<16xi32> + %ptrs = "tt.addptr"(%base, %indices) + : (!tt.ptr, tensor<16xi32>) -> tensor<16x!tt.ptr> + %values = "tt.load"(%ptrs) + : (tensor<16x!tt.ptr>) -> tensor<16xf16> + %reduced = "tt.reduce"(%values) : (tensor<16xf16>) -> f16 + "tt.store"(%base, %reduced) : (!tt.ptr, f16) -> () + return + } + } + )mlir", + &context); + ASSERT_TRUE(module); + + llvm::SmallVector roots; + module->walk([&](mlir::func::FuncOp function) { + for (mlir::Operation &operation : function.getBody().front()) + if (!operation.hasTrait()) + roots.push_back(&operation); + }); + ASSERT_EQ(roots.size(), 5u); + + mlir::ascend::SimtAnchorDescriptor anchor; + anchor.operation = roots[1]; + anchor.scopeOperations = {roots[1]}; + anchor.scopeInsertionPoint = roots[1]; + anchor.kind = mlir::ascend::SimtAnchorKind::DirectGather; + anchor.lowerability.mixed = true; + anchor.materializable = true; + mlir::ascend::SimtAnchorPlan anchorPlan; + anchorPlan.anchors.push_back(std::move(anchor)); + + mlir::ascend::PhaseBoundaryPlan phasePlan{ + mlir::ascend::PhaseBoundaryDomain::LoadedIndexRowwiseReduction, + "loaded_index_rowwise_reduction", std::nullopt}; + phasePlan.rootOperations.assign(roots.begin(), roots.end()); + phasePlan.rootPhaseIds = {"row_dispatch", "row_load", "row_load", + "row_reduction", "convert_store"}; + auto result = mlir::ascend::StageBoundaryAnalysis().analyze( + phasePlan, SimdSimtFeatureSummary{}, &anchorPlan); + if (!result) + FAIL() << llvm::toString(result.takeError()); + + const LogicalStage *gather = nullptr; + for (const LogicalPhase &phase : result->phases) + for (const LogicalStage &stage : phase.stages) + if (stage.id == "indirect_row_gather") + gather = &stage; + ASSERT_NE(gather, nullptr); + EXPECT_FALSE(gather->localSimtMaterializable); + EXPECT_TRUE(gather->localSimtFactors.empty()); + EXPECT_TRUE(gather->simtAnchorIndices.empty()); +} + TEST(SimdSimtCostModelTest, PointerInductionLoopIsNotADataRecurrence) { mlir::MLIRContext context; context.getOrLoadDialect(); @@ -1243,14 +1183,3 @@ TEST(SimdSimtCostModelTest, IncompatibleDominantStructuresRequireStageSplit) { EXPECT_NE(llvm::toString(std::move(error)).find("requires_split"), std::string::npos); } - -TEST(SimdSimtCostModelTest, FeatureSummaryPartitionIsExplicitFallback) { - auto result = StagePartitioner().partition(triangularBt16StageFeatures(), - StagePartitionerOptions{}); - if (!result) - FAIL() << llvm::toString(result.takeError()); - ASSERT_TRUE(*result); - EXPECT_EQ((**result).boundarySource, "feature_summary_fallback"); - EXPECT_FALSE((**result).operationOwnershipComplete); - EXPECT_EQ((**result).modeledOperationCount, 0); -} diff --git a/third_party/ascend/unittest/costmodel_ut/test_compiler_costmodel_contract.py b/third_party/ascend/unittest/costmodel_ut/test_compiler_costmodel_contract.py index d131c63ece..606952030b 100644 --- a/third_party/ascend/unittest/costmodel_ut/test_compiler_costmodel_contract.py +++ b/third_party/ascend/unittest/costmodel_ut/test_compiler_costmodel_contract.py @@ -10,6 +10,20 @@ class CompilerCostmodelContractTest(unittest.TestCase): @staticmethod def _load_compiler_module(): + stubbed_modules = [ + "ctypes", + "triton", + "triton._C", + "triton._C.libtriton", + "triton._C.libtriton.ascend", + "triton.backends.ascend", + "triton.backends.ascend.utils", + "triton.backends.ascend.driver", + "triton.backends.compiler", + "triton.runtime", + "triton.runtime.cache", + ] + saved_modules = {name: sys.modules.get(name) for name in stubbed_modules} ctypes_stub = types.ModuleType("ctypes") ctypes_stub.c_int64 = int sys.modules["ctypes"] = ctypes_stub @@ -123,6 +137,11 @@ def put(self, payload, file_name, binary=False): module = importlib.util.module_from_spec(spec) assert spec and spec.loader spec.loader.exec_module(module) + for name, saved in saved_modules.items(): + if saved is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = saved return module, dump_mgr, GPUTarget def test_parse_options_costmodel_forces_no_bytecode(self): @@ -187,13 +206,22 @@ def test_all_simd_decision_replaces_route_request_with_backend_mode(self): def test_mixed_decision_preserves_route_request_and_factor(self): cmplr, _dump_mgr, _GPUTarget = self._load_compiler_module() - metadata = {"compile_mode": "simd_simt"} + metadata = { + "compile_mode": "simd_simt", + "route_transform_v1_materializable": True, + } - cmplr._apply_cpp_simd_simt_decision(metadata, "mixed_simd_simt", 4, "{\"decision\":\"mixed\"}") + cmplr._apply_cpp_simd_simt_decision(metadata, "mixed_simd_simt", 4, "{\"decision\":\"mixed\"}", 2) self.assertEqual(metadata["compile_mode"], "simd_simt") self.assertEqual(metadata["auto_simt_superblock_factor"], 4) self.assertEqual(metadata["auto_simt_requested_kind"], "mixed_simd_simt") + self.assertTrue(metadata["auto_blockify_v1_enabled"]) + self.assertTrue(metadata["auto_blockify_v1_runtime_cap"]) + self.assertEqual(metadata["auto_simt_scope_superblock_factor"], 4) + self.assertNotIn("auto_simt_scope_num_warps", metadata) + self.assertNotIn("num_warps", metadata) + self.assertNotIn("scope_superblock_backend_abi_version", metadata) if __name__ == "__main__": diff --git a/third_party/ascend/unittest/pytest_ut/test_auto_blockify_v1_policy.py b/third_party/ascend/unittest/pytest_ut/test_auto_blockify_v1_policy.py index 607ae85f0e..509d71b7b9 100644 --- a/third_party/ascend/unittest/pytest_ut/test_auto_blockify_v1_policy.py +++ b/third_party/ascend/unittest/pytest_ut/test_auto_blockify_v1_policy.py @@ -23,10 +23,18 @@ import pytest -from triton.backends.ascend.compiler import _resolve_auto_blockify_v1_policy +from triton.backends.ascend.compiler import ( + _build_costmodel_analysis_ttir, + _can_materialize_scope_superblock, + _publish_route_transform_capability, + _resolve_auto_blockify_v1_policy, + _selected_npuir_superblock_factor, +) SAFE_TTIR = "module { tt.func public @safe() { tt.return } }" ATOMIC_TTIR = "module { tt.func public @atomic() { %0 = tt.atomic_rmw add } }" +CACHE_MODIFIER_TTIR = ( + 'module { tt.func public @cached() { tt.store %ptr, %value {cacheModifier = 1 : i32} : tensor<8x!tt.ptr> } }') @pytest.mark.parametrize( @@ -62,3 +70,89 @@ def test_auto_blockify_v1_explicit_blacklist_override_is_preserved(): opt = SimpleNamespace(enable_auto_blockify=True) with patch("triton.backends.ascend.compiler._is_auto_map_parallel_blocks_enabled", return_value=False): assert _resolve_auto_blockify_v1_policy(ATOMIC_TTIR, metadata, opt) + + +def test_cache_modifier_does_not_disable_auto_blockify_v1(): + metadata = {} + opt = SimpleNamespace(enable_auto_blockify=True) + with patch("triton.backends.ascend.compiler._is_auto_map_parallel_blocks_enabled", return_value=False): + assert _resolve_auto_blockify_v1_policy(CACHE_MODIFIER_TTIR, metadata, opt) + assert not metadata["has_auto_blockify_blacklist_op"] + + +def test_route_transform_capability_is_single_resolved_fact(): + metadata = { + "ttir_layout_merge_applied": True, + "ttir_layout_coalesce_factor": 8, + "ttir_layout_coalesce_axis": 0, + "auto_blockify_v1_requested": True, + "auto_blockify_v1_enabled": True, + "auto_blockify_v1_disable_reasons": [], + } + opt = SimpleNamespace(compile_on_910_95=True, num_warps=4, logical_program_count_hint=9) + capability = __import__("json").loads(_publish_route_transform_capability(metadata, opt)) + assert capability["row_coalescing_applied"] + assert capability["row_coalescing_factor"] == 8 + assert capability["auto_blockify_v1_materializable"] + assert capability["whole_kernel_superblock_factors"] == [1, 2, 4] + assert capability["scope_superblock_factors"] == [1, 2, 4] + assert capability["logical_program_count_hint"] == 9 + assert capability["superblock_runtime_groups"]["4"] == { + "full_group_count": 2, + "tail_count": 1, + } + + +@pytest.mark.parametrize( + "compile_on_910_95,num_warps,v1_materializable,expected", + [ + (True, 1, True, True), + (True, 4, True, True), + (True, 4, False, False), + (False, 4, True, False), + (True, 0, True, False), + ], +) +def test_scope_superblock_uses_npuir_abi_v2(compile_on_910_95, num_warps, v1_materializable, expected): + opt = SimpleNamespace(compile_on_910_95=compile_on_910_95, num_warps=num_warps) + assert _can_materialize_scope_superblock({}, opt, v1_materializable) is expected + + +@pytest.mark.parametrize( + "metadata,option_factor,expected", + [ + ({}, 2, 2), + ({"auto_simt_scope_superblock_factor": 1}, 4, 4), + ({ + "auto_simt_effective_kind": "mixed_simd_simt", + "auto_simt_scope_superblock_factor": 2, + }, 4, 1), + ({ + "compile_mode": "simd_simt", + "auto_simt_scope_superblock_factor": 4, + }, 4, 1), + ({ + "auto_simt_effective_kind": "all_simt_only", + "auto_simt_superblock_factor": 4, + }, 2, 4), + ], +) +def test_selected_npuir_superblock_factor_respects_route_owner(metadata, option_factor, expected): + opt = SimpleNamespace(superblock_factor=option_factor) + assert _selected_npuir_superblock_factor(metadata, opt) == expected + + +def test_costmodel_analysis_view_materializes_v1_only_on_clone(): + metadata = {"auto_blockify_v1_enabled": True} + original = SimpleNamespace(context=object()) + analysis = SimpleNamespace() + with patch("triton.backends.ascend.compiler._parse_ttir_text", return_value=analysis) as parse, \ + patch("triton.backends.ascend.compiler._run_ta_simt_auto_blockify_v1", return_value=True) as run_v1: + result = _build_costmodel_analysis_ttir(original, metadata, SimpleNamespace()) + parse.assert_called_once() + run_v1.assert_called_once() + assert run_v1.call_args.args[0] is analysis + assert run_v1.call_args.kwargs["super_block_factor"] == 1 + assert metadata["auto_simt_costmodel_analysis_v1_materialized"] + assert metadata["auto_simt_costmodel_analysis_ir"] == "post_auto_blockify_v1_f1_ttir" + assert result == str(analysis) diff --git a/third_party/ascend/unittest/pytest_ut/test_simd_simt_costmodel_cases.py b/third_party/ascend/unittest/pytest_ut/test_simd_simt_costmodel_cases.py new file mode 100644 index 0000000000..3dbc1c3c4e --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_simd_simt_costmodel_cases.py @@ -0,0 +1,440 @@ +import csv +import json +import os +import statistics + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl +import triton.runtime.driver as driver +from triton.backends.ascend.utils import is_compile_on_910_95 + +simd_simt_910_95_only = pytest.mark.xfail( + not is_compile_on_910_95(), + reason="SIMD/SIMT cost model only supports 910_95", + run=False, +) + + +def _vector_core_count(): + properties = driver.active.utils.get_device_properties(torch.npu.current_device()) + return int(properties["num_vectorcore"]) + + +def _load_route_report(path, expected): + report = json.loads(path.read_text()) + assert report["stage_model"]["applied"] + assert report["effective_decision_kind"] == expected + return report + + +def _launch_options(report_path, logical_programs): + options = { + "num_warps": 4, + "compile_mode": "simd_simt", + "auto_simt_scope_mode": "auto", + "auto_simt_scope_dump": str(report_path), + "logical_program_count_hint": logical_programs, + "physical_vector_core_count_hint": _vector_core_count(), + } + if os.getenv("TRITON_TEST_DISABLE_TTIR_LAYOUT_MERGE") == "1": + options["enable_ttir_layout_merge"] = False + return options + + +def _assert_performance(case, launch, profile_root, documented_us, tolerance=1.2): + for _ in range(20): + launch() + torch.npu.synchronize() + config = torch_npu.profiler._ExperimentalConfig( + aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization, + profiler_level=torch_npu.profiler.ProfilerLevel.Level1, + l2_cache=False, + data_simplification=False, + ) + skip_first, warmup, active = 5, 3, 20 + with torch_npu.profiler.profile( + activities=[torch_npu.profiler.ProfilerActivity.NPU], + schedule=torch_npu.profiler.schedule( + wait=0, + warmup=warmup, + active=active, + repeat=1, + skip_first=skip_first, + ), + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(str(profile_root)), + record_shapes=False, + profile_memory=False, + with_stack=False, + with_flops=False, + with_modules=False, + experimental_config=config, + ) as profiler: + for _ in range(skip_first + warmup + active): + launch() + profiler.step() + torch.npu.synchronize() + + detail_files = list(profile_root.rglob("kernel_details.csv")) + assert detail_files, f"{case}: profiler did not generate kernel_details.csv" + durations = [] + for detail_file in detail_files: + with detail_file.open(newline="") as stream: + for row in csv.DictReader(stream): + if row.get("Duration(us)"): + durations.append(float(row["Duration(us)"])) + assert durations, f"{case}: profiler generated no kernel duration" + duration_us = statistics.median(durations) + maximum_us = documented_us * tolerance + print( + f"{case}: profiler median {duration_us:.3f} us (documented {documented_us:.3f} us, limit {maximum_us:.3f} us)") + assert duration_us <= maximum_us + + +@triton.jit +def gather_dot_min( + a_ptr, + b_ptr, + indices_ptr, + out_ptr, + M, + N, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_om, + stride_on, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_k = tl.arange(0, BLOCK_K) + gather_k = tl.load(indices_ptr + offs_k) + a = tl.load(a_ptr + offs_m[:, None] * stride_am + gather_k[None, :] * stride_ak) + b = tl.load(b_ptr + gather_k[:, None] * stride_bk + offs_n[None, :] * stride_bn) + result = tl.dot(a, b) + tl.store(out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on, result) + + +@simd_simt_910_95_only +def test_costmodel_gather_dot_min(tmp_path): + logical_programs = _vector_core_count() + block = 16 + source_k = 256 + a = torch.randn((logical_programs * block, source_k), dtype=torch.float16, device="npu") + b = torch.randn((source_k, block), dtype=torch.float16, device="npu") + indices = torch.tensor( + [10, 25, 100, 200, 5, 50, 150, 255, 1, 2, 3, 4, 6, 7, 8, 9], + dtype=torch.int32, + device="npu", + ) + output = torch.empty((logical_programs * block, block), dtype=torch.float32, device="npu") + report_path = tmp_path / "gather_dot_min_route.json" + + def launch(): + gather_dot_min[(logical_programs, 1)]( + a, + b, + indices, + output, + a.shape[0], + b.shape[1], + a.stride(0), + a.stride(1), + b.stride(0), + b.stride(1), + output.stride(0), + output.stride(1), + BLOCK_M=block, + BLOCK_N=block, + BLOCK_K=block, + **_launch_options(report_path, logical_programs), + ) + + launch() + expected = torch.matmul(a[:, indices].float(), b[indices, :].float()) + torch.testing.assert_close(output, expected, rtol=1e-2, atol=1e-2) + report = _load_route_report(report_path, "all_simt_only") + assert report["features"]["dot_ops"] == 1 + _assert_performance("gather_dot_min", launch, tmp_path / "gather_profile", 5.478, tolerance=1.35) + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.jit(do_not_specialize=["T"]) +def merge_16x16_to_64x64_inverse_kernel( + A, + Ai, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + USE_TMA: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + i_t = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos = tl.load(cu_seqlens + i_n).to(tl.int32) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, 16) + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + A += (bos * H + i_h) * BT + Ai += (bos * H + i_h) * BT + + p_A_11 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) + p_A_22 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) + p_A_33 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0)) + p_A_44 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0)) + b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) + b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) + b_Ai_33 = tl.load(p_A_33, boundary_check=(0, 1)).to(tl.float32) + b_Ai_44 = tl.load(p_A_44, boundary_check=(0, 1)).to(tl.float32) + b_Ai_11 = -tl.where(m_A, b_Ai_11, 0.0) + b_Ai_22 = -tl.where(m_A, b_Ai_22, 0.0) + b_Ai_33 = -tl.where(m_A, b_Ai_33, 0.0) + b_Ai_44 = -tl.where(m_A, b_Ai_44, 0.0) + + for i in range(2, min(16, T - i_t * BT)): + b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i) + b_a_11 = tl.where(o_i < i, b_a_11, 0.0) + b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0) + b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11) + for i in range(18, min(32, T - i_t * BT)): + b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16) + b_a_22 = tl.where(o_i < i - 16, b_a_22, 0.0) + b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0) + b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22) + for i in range(34, min(48, T - i_t * BT)): + b_a_33 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 32) + b_a_33 = tl.where(o_i < i - 32, b_a_33, 0.0) + b_a_33 += tl.sum(b_a_33[:, None] * b_Ai_33, 0) + b_Ai_33 = tl.where((o_i == i - 32)[:, None], b_a_33, b_Ai_33) + for i in range(50, min(64, T - i_t * BT)): + b_a_44 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 48) + b_a_44 = tl.where(o_i < i - 48, b_a_44, 0.0) + b_a_44 += tl.sum(b_a_44[:, None] * b_Ai_44, 0) + b_Ai_44 = tl.where((o_i == i - 48)[:, None], b_a_44, b_Ai_44) + b_Ai_11 += m_I + b_Ai_22 += m_I + b_Ai_33 += m_I + b_Ai_44 += m_I + + p_A_21 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) + p_A_31 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0)) + p_A_32 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0)) + p_A_41 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0)) + p_A_42 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0)) + p_A_43 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0)) + b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) + b_A_31 = tl.load(p_A_31, boundary_check=(0, 1)).to(tl.float32) + b_A_32 = tl.load(p_A_32, boundary_check=(0, 1)).to(tl.float32) + b_A_41 = tl.load(p_A_41, boundary_check=(0, 1)).to(tl.float32) + b_A_42 = tl.load(p_A_42, boundary_check=(0, 1)).to(tl.float32) + b_A_43 = tl.load(p_A_43, boundary_check=(0, 1)).to(tl.float32) + + b_Ai_21 = -tl.dot(tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION), b_Ai_11, input_precision=DOT_PRECISION) + b_Ai_32 = -tl.dot(tl.dot(b_Ai_33, b_A_32, input_precision=DOT_PRECISION), b_Ai_22, input_precision=DOT_PRECISION) + b_Ai_43 = -tl.dot(tl.dot(b_Ai_44, b_A_43, input_precision=DOT_PRECISION), b_Ai_33, input_precision=DOT_PRECISION) + b_Ai_31 = -tl.dot( + b_Ai_33, + tl.dot(b_A_31, b_Ai_11, input_precision=DOT_PRECISION) + tl.dot(b_A_32, b_Ai_21, input_precision=DOT_PRECISION), + input_precision=DOT_PRECISION, + ) + b_Ai_42 = -tl.dot( + b_Ai_44, + tl.dot(b_A_42, b_Ai_22, input_precision=DOT_PRECISION) + tl.dot(b_A_43, b_Ai_32, input_precision=DOT_PRECISION), + input_precision=DOT_PRECISION, + ) + b_Ai_41 = -tl.dot( + b_Ai_44, + tl.dot(b_A_41, b_Ai_11, input_precision=DOT_PRECISION) + + tl.dot(b_A_42, b_Ai_21, input_precision=DOT_PRECISION) + tl.dot(b_A_43, b_Ai_31, input_precision=DOT_PRECISION), + input_precision=DOT_PRECISION, + ) + + p_Ai_11 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) + p_Ai_22 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) + p_Ai_33 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0)) + p_Ai_44 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0)) + p_Ai_21 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) + p_Ai_31 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0)) + p_Ai_32 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0)) + p_Ai_41 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0)) + p_Ai_42 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0)) + p_Ai_43 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0)) + tl.store(p_Ai_11, b_Ai_11, boundary_check=(0, 1)) + tl.store(p_Ai_22, b_Ai_22, boundary_check=(0, 1)) + tl.store(p_Ai_33, b_Ai_33, boundary_check=(0, 1)) + tl.store(p_Ai_44, b_Ai_44, boundary_check=(0, 1)) + tl.store(p_Ai_21, b_Ai_21, boundary_check=(0, 1)) + tl.store(p_Ai_31, b_Ai_31, boundary_check=(0, 1)) + tl.store(p_Ai_32, b_Ai_32, boundary_check=(0, 1)) + tl.store(p_Ai_41, b_Ai_41, boundary_check=(0, 1)) + tl.store(p_Ai_42, b_Ai_42, boundary_check=(0, 1)) + tl.store(p_Ai_43, b_Ai_43, boundary_check=(0, 1)) + + +@simd_simt_910_95_only +def test_costmodel_solve_tril(tmp_path): + batch, sequence_length, heads, block = 1, 1024, 32, 64 + chunks = sequence_length // block + logical_programs = batch * chunks * heads + torch.manual_seed(1) + lower = torch.tril(torch.randn((block, block), dtype=torch.float32), diagonal=-1) * 0.01 + source = lower.reshape(1, block, 1, block).expand(chunks, block, heads, + block).reshape(batch, sequence_length, heads, block) + a = source.to(device="npu", dtype=torch.float32) + output = torch.zeros((batch, sequence_length, heads, block), dtype=torch.float32, device="npu") + report_path = tmp_path / "solve_tril_route.json" + + def launch(): + merge_16x16_to_64x64_inverse_kernel[(chunks, batch * heads)]( + a, + output, + None, + None, + sequence_length, + H=heads, + BT=block, + USE_TMA=False, + DOT_PRECISION="ieee", + **_launch_options(report_path, logical_programs), + ) + + launch() + inverse = torch.linalg.inv(torch.eye(block) + lower) + expected = inverse.reshape(1, block, 1, block).expand(chunks, block, heads, block).reshape_as(output).to("npu") + torch.testing.assert_close(output, expected, rtol=3e-2, atol=3e-2) + report = _load_route_report(report_path, "mixed_simd_simt") + assert report["materialized_simt_anchor_count"] > 0 + _assert_performance("solve_tril", launch, tmp_path / "solve_profile", 197.198) + + +@triton.jit +def _fbgemm_gather_scale_fp8_rowwise_quant_dense_tokens( + output_ptr, + output_scale_ptr, + input_ptr, + token_indices_ptr, + expert_indices_ptr, + scores_ptr, + scale_ub_ptr, + stride_t, + stride_e, + valid_token_count, + D: tl.constexpr, + TL_FP8_DTYPE: tl.constexpr, + MAX_FP8: tl.constexpr, + EPS: tl.constexpr, + CLAMP_MAX: tl.constexpr, + BLOCK_D: tl.constexpr, +): + tl.static_assert(D % BLOCK_D == 0, "D must be a multiple of BLOCK_D") + output_token = tl.program_id(0) + valid_token_count = tl.load(valid_token_count, None, eviction_policy="evict_last") + if output_token >= valid_token_count: + return + input_token = tl.load(token_indices_ptr + output_token) + expert = tl.load(expert_indices_ptr + output_token) + score = tl.load(scores_ptr + input_token * stride_t + expert * stride_e).to(tl.float32) + offsets = tl.arange(0, BLOCK_D) + input_block = input_ptr + input_token.to(tl.int64) * D + offsets + row_max = 0.0 + for _ in range(0, D, BLOCK_D): + values = tl.load(input_block, eviction_policy="evict_last").to(tl.float32) * score + row_max = tl.maximum(tl.max(tl.abs(values)), row_max) + input_block += BLOCK_D + + if CLAMP_MAX: + row_max = tl.clamp(row_max, EPS, tl.load(scale_ub_ptr)) + else: + row_max = tl.maximum(row_max, EPS) + scale = MAX_FP8 / row_max + tl.store(output_scale_ptr + output_token, 1.0 / scale) + input_block = input_ptr + input_token.to(tl.int64) * D + offsets + output_block = output_ptr + output_token.to(tl.int64) * D + offsets + for _ in range(0, D, BLOCK_D): + values = tl.load(input_block, eviction_policy="evict_first").to(tl.float32) * score + quantized = tl.clamp(values * scale, -MAX_FP8, MAX_FP8).to(TL_FP8_DTYPE) + tl.store(output_block, quantized, cache_modifier=".cg") + input_block += BLOCK_D + output_block += BLOCK_D + + +@simd_simt_910_95_only +def test_costmodel_fbgemm_rowwise_quant(tmp_path): + tokens, width, experts, valid = 256, 1024, 8, 512 + torch.manual_seed(2) + input_tensor = torch.randn((tokens, width), dtype=torch.float16, device="npu") + token_indices = torch.arange(valid, dtype=torch.int32, device="npu") % tokens + expert_indices = torch.arange(valid, dtype=torch.int32, device="npu") % experts + scores = torch.randn((tokens, experts), dtype=torch.float16, device="npu") + valid_count = torch.tensor([valid], dtype=torch.int32, device="npu") + scale_ub = torch.tensor([448.0], dtype=torch.float32, device="npu") + output = torch.empty((valid, width), dtype=torch.float8_e4m3fn, device="npu") + output_scale = torch.empty((valid, ), dtype=torch.float32, device="npu") + report_path = tmp_path / "fbgemm_route.json" + + def launch(): + _fbgemm_gather_scale_fp8_rowwise_quant_dense_tokens[(valid, )]( + output, + output_scale, + input_tensor, + token_indices, + expert_indices, + scores, + scale_ub, + scores.stride(0), + scores.stride(1), + valid_count, + D=width, + TL_FP8_DTYPE=tl.float8e4nv, + MAX_FP8=448.0, + EPS=1.0e-12, + CLAMP_MAX=False, + BLOCK_D=width, + **_launch_options(report_path, valid), + ) + + launch() + gathered = input_tensor[token_indices].float() * scores[token_indices, expert_indices].float()[:, None] + row_max = torch.clamp(torch.amax(torch.abs(gathered), dim=1), min=1.0e-12) + expected_scale = row_max / 448.0 + expected = torch.clamp(gathered / expected_scale[:, None], -448.0, 448.0).to(torch.float8_e4m3fn) + torch.testing.assert_close(output_scale, expected_scale, rtol=2e-3, atol=2e-3) + torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=16) + layout_merge_disabled = os.getenv("TRITON_TEST_DISABLE_TTIR_LAYOUT_MERGE") == "1" + expected_route = "all_simd" if layout_merge_disabled else "all_simt_only" + report = _load_route_report(report_path, expected_route) + capability = report["route_transform_capability"] + assert capability["source_logical_program_count_hint"] == valid + if layout_merge_disabled: + assert not capability["row_coalescing_applied"] + assert capability["logical_program_count_hint"] == valid + else: + assert capability["row_coalescing_factor"] == 2 + assert capability["logical_program_count_hint"] == valid // 2 + documented_us = 20.578 if layout_merge_disabled else 8.904 + _assert_performance( + "fbgemm_rowwise_quant", + launch, + tmp_path / "fbgemm_profile", + documented_us, + tolerance=1.25, + ) From 463f349f06d513837a28e75c32963610270c37fb Mon Sep 17 00:00:00 2001 From: yangkaixin Date: Thu, 27 Aug 2026 20:15:13 +0800 Subject: [PATCH 2/5] fixed delayed bug on inserting wrong sync flag for simt scope with >1 super blocks --- third_party/ascend/backend/compiler.py | 8 -------- .../costmodel_ut/test_compiler_costmodel_contract.py | 7 +++++++ .../unittest/pytest_ut/test_simd_simt_costmodel_cases.py | 4 ++++ 3 files changed, 11 insertions(+), 8 deletions(-) diff --git a/third_party/ascend/backend/compiler.py b/third_party/ascend/backend/compiler.py index 8e7c53bf7d..309d3f5543 100644 --- a/third_party/ascend/backend/compiler.py +++ b/third_party/ascend/backend/compiler.py @@ -1254,14 +1254,6 @@ def linalg_to_bin_enable_npu_compile_910_95(linalg: str, metadata, opt): _compile_option_list += [f"--num-warps={num_warps}"] warp_size = metadata.get("warp_size") or opt.warp_size _compile_option_list += [f"--threads-per-warp={warp_size}"] - # A result-bearing local SIMT scope creates explicit AIV -> AIC - # values at the scope boundary. The delayed solver currently - # cannot represent that edge and aborts while walking between - # anchors; the regular cross-core solver handles the same IR. - # Key this off the lowered IR contract so explicit and - # cost-model-materialized scopes follow the same path. - _compile_option_list += ["--enable-hivm-delayed-cross-core-gss=false"] - partition_mode = _validate_partition_and_bind_sub_block( metadata.get( "enable_partition_and_bind_sub_block", diff --git a/third_party/ascend/unittest/costmodel_ut/test_compiler_costmodel_contract.py b/third_party/ascend/unittest/costmodel_ut/test_compiler_costmodel_contract.py index 606952030b..2523d708e6 100644 --- a/third_party/ascend/unittest/costmodel_ut/test_compiler_costmodel_contract.py +++ b/third_party/ascend/unittest/costmodel_ut/test_compiler_costmodel_contract.py @@ -1,4 +1,5 @@ import importlib.util +import inspect import sys import tempfile import types @@ -223,6 +224,12 @@ def test_mixed_decision_preserves_route_request_and_factor(self): self.assertNotIn("num_warps", metadata) self.assertNotIn("scope_superblock_backend_abi_version", metadata) + def test_mixed_compile_keeps_delayed_cross_core_gss_enabled(self): + cmplr, _dump_mgr, _GPUTarget = self._load_compiler_module() + + source = inspect.getsource(cmplr.linalg_to_bin_enable_npu_compile_910_95) + self.assertNotIn("--enable-hivm-delayed-cross-core-gss=false", source) + if __name__ == "__main__": unittest.main() diff --git a/third_party/ascend/unittest/pytest_ut/test_simd_simt_costmodel_cases.py b/third_party/ascend/unittest/pytest_ut/test_simd_simt_costmodel_cases.py index 3dbc1c3c4e..c8bbe5c235 100644 --- a/third_party/ascend/unittest/pytest_ut/test_simd_simt_costmodel_cases.py +++ b/third_party/ascend/unittest/pytest_ut/test_simd_simt_costmodel_cases.py @@ -323,6 +323,10 @@ def launch(): torch.testing.assert_close(output, expected, rtol=3e-2, atol=3e-2) report = _load_route_report(report_path, "mixed_simd_simt") assert report["materialized_simt_anchor_count"] > 0 + assert report["selected_superblock_factor"] == 4 + assert report["effective_runtime_factor"] == 4 + assert report["full_group_count"] == logical_programs // 4 + assert report["tail_count"] == 0 _assert_performance("solve_tril", launch, tmp_path / "solve_profile", 197.198) From 6debb4e3fad00bd66c154b2fa48689abc191e392 Mon Sep 17 00:00:00 2001 From: yangkaixin Date: Fri, 28 Aug 2026 15:25:28 +0800 Subject: [PATCH 3/5] local simt scope support sublocks in CV mixed case --- .../ascend/patch/triton-ascend-3.6.0.patch | 8 +-- .../General/TritonToLinalg/simt_scope.mlir | 4 ++ .../pytest_ut/test_auto_blockify_v1_policy.py | 7 +-- ...indirect_load_pointer_cast_precise_size.py | 3 +- .../test_simd_simt_costmodel_cases.py | 49 +++++++++++++++---- 5 files changed, 54 insertions(+), 17 deletions(-) diff --git a/third_party/ascend/patch/triton-ascend-3.6.0.patch b/third_party/ascend/patch/triton-ascend-3.6.0.patch index 8a5a4e8d84..1e24bfa9d8 100644 --- a/third_party/ascend/patch/triton-ascend-3.6.0.patch +++ b/third_party/ascend/patch/triton-ascend-3.6.0.patch @@ -227,7 +227,7 @@ diff --git a/bin/RegisterTritonDialects.h b/bin/RegisterTritonDialects.h index 48049f026..e96284bc2 100644 --- a/bin/RegisterTritonDialects.h +++ b/bin/RegisterTritonDialects.h -@@ -1,6 +1,36 @@ +@@ -1,6 +1,37 @@ #pragma once +#include "ascend/include/AutoBlockify/Passes.h" +#include "ascend/include/AutoBlockifyV1/Passes.h" @@ -243,6 +243,7 @@ index 48049f026..e96284bc2 100644 +#include "ascend/include/DynamicCVPipeline/SplitDataflow/RefineArgsBlockId.h" +#include "ascend/include/DynamicCVPipeline/StandardizeOp.h" +#include "ascend/include/TritonControlFlowOpt/Passes.h" ++#include "ascend/include/TTIRLayoutMerge/Passes.h" +#include "ascend/include/TritonToAnnotation/Passes.h" +#include "ascend/include/TritonToGraph/Passes.h" +#include "ascend/include/TritonToHFusion/Passes.h" @@ -264,7 +265,7 @@ index 48049f026..e96284bc2 100644 #include "nvidia/include/Dialect/NVGPU/IR/Dialect.h" #include "nvidia/include/Dialect/NVWS/IR/Dialect.h" #include "proton/Dialect/include/Conversion/ProtonGPUToLLVM/Passes.h" -@@ -85,6 +115,20 @@ inline void registerTritonDialects(mlir::DialectRegistry ®istry) { +@@ -85,6 +116,21 @@ inline void registerTritonDialects(mlir::DialectRegistry ®istry) { mlir::triton::registerConvertTritonGPUToLLVMPass(); mlir::triton::registerConvertNVGPUToLLVMPass(); mlir::triton::registerAllocateSharedMemoryNvPass(); @@ -279,13 +280,14 @@ index 48049f026..e96284bc2 100644 + mlir::triton::registerTritonToLLVMPasses(); + mlir::triton::registerAutoBlockifyPasses(); + mlir::triton::registerAutoBlockifyV1Passes(); ++ mlir::triton::registerTTIRLayoutMergePasses(); + mlir::triton::cfg::registerTritonToGraphPasses(); + mlir::triton::registerBubbleUpOperationPass(); + mlir::registerLLVMDIScope(); mlir::LLVM::registerInlinerInterface(registry); mlir::NVVM::registerInlinerInterface(registry); -@@ -136,16 +180,38 @@ inline void registerTritonDialects(mlir::DialectRegistry ®istry) { +@@ -136,16 +182,38 @@ inline void registerTritonDialects(mlir::DialectRegistry ®istry) { mlir::triton::proton::gpu::registerScheduleBufferStorePass(); mlir::triton::proton::gpu::registerAddSchedBarriersPass(); diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/simt_scope.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/simt_scope.mlir index 36773da28b..e86daa3bc8 100644 --- a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/simt_scope.mlir +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/simt_scope.mlir @@ -5,13 +5,17 @@ // CHECK-LABEL: func.func @simt_scope // CHECK-SAME: parallel_mode = "mix_simd_simt" // CHECK: scope.scope +// CHECK: } { // CHECK-SAME: hivm.func_core_type = #hivm.func_core_type // CHECK-SAME: hivm.vf_mode = #hivm.vf_mode // CHECK-SAME: no_inline // CHECK-SAME: outline // CHECK-SAME: vector_mode = "simt" tt.func public @simt_scope(%arg0: !tt.ptr) { + %zero = arith.constant dense<0.000000e+00> : tensor<16xf32> + %ptrs = tt.splat %arg0 : !tt.ptr -> tensor<16x!tt.ptr> scope.scope : () -> () { + tt.store %ptrs, %zero : tensor<16x!tt.ptr> scope.return } {vector_mode = "simt"} tt.return diff --git a/third_party/ascend/unittest/pytest_ut/test_auto_blockify_v1_policy.py b/third_party/ascend/unittest/pytest_ut/test_auto_blockify_v1_policy.py index 509d71b7b9..7d1eeced77 100644 --- a/third_party/ascend/unittest/pytest_ut/test_auto_blockify_v1_policy.py +++ b/third_party/ascend/unittest/pytest_ut/test_auto_blockify_v1_policy.py @@ -96,10 +96,11 @@ def test_route_transform_capability_is_single_resolved_fact(): assert capability["auto_blockify_v1_materializable"] assert capability["whole_kernel_superblock_factors"] == [1, 2, 4] assert capability["scope_superblock_factors"] == [1, 2, 4] - assert capability["logical_program_count_hint"] == 9 + assert capability["source_logical_program_count_hint"] == 9 + assert capability["logical_program_count_hint"] == 2 assert capability["superblock_runtime_groups"]["4"] == { - "full_group_count": 2, - "tail_count": 1, + "full_group_count": 0, + "tail_count": 2, } diff --git a/third_party/ascend/unittest/pytest_ut/test_indirect_load_pointer_cast_precise_size.py b/third_party/ascend/unittest/pytest_ut/test_indirect_load_pointer_cast_precise_size.py index 0f8d5ae48a..96e8d09b90 100644 --- a/third_party/ascend/unittest/pytest_ut/test_indirect_load_pointer_cast_precise_size.py +++ b/third_party/ascend/unittest/pytest_ut/test_indirect_load_pointer_cast_precise_size.py @@ -32,10 +32,11 @@ def test_indirect_load_pointer_cast_precise_size_e2e(): block_size = 4 block = 8 - block_table_cpu = torch.tensor([10, 11, 12, 13], dtype=torch.int32) + block_table_cpu = torch.tensor([10, 11, 12, 13], dtype=torch.int32, device="cpu") positions_cpu = torch.tensor( [0, 1, 4, 7, 8, 11, 12, 15], dtype=torch.int32, + device="cpu", ) block_table = block_table_cpu.npu() diff --git a/third_party/ascend/unittest/pytest_ut/test_simd_simt_costmodel_cases.py b/third_party/ascend/unittest/pytest_ut/test_simd_simt_costmodel_cases.py index c8bbe5c235..4d4cd107fe 100644 --- a/third_party/ascend/unittest/pytest_ut/test_simd_simt_costmodel_cases.py +++ b/third_party/ascend/unittest/pytest_ut/test_simd_simt_costmodel_cases.py @@ -290,17 +290,32 @@ def merge_16x16_to_64x64_inverse_kernel( tl.store(p_Ai_43, b_Ai_43, boundary_check=(0, 1)) +SOLVE_TRIL_CASES = [ + (1, 1024, 32, 64, 182.929), + (1, 1024, 64, 64, 613.346), + (8, 1024, 32, 64, 1396.827), + (16, 1024, 64, 64, 5474.364), + (1, 8192, 64, 64, 2730.823), + (8, 8192, 32, 64, 10930.605), + (16, 8192, 64, 64, 43821.177), + (4, 131072, 32, 64, 87436.743), +] + + @simd_simt_910_95_only -def test_costmodel_solve_tril(tmp_path): - batch, sequence_length, heads, block = 1, 1024, 32, 64 +@pytest.mark.parametrize( + "batch,sequence_length,heads,block,documented_us", + SOLVE_TRIL_CASES, + ids=[f"B{b}-T{t}-H{h}-BT{bt}" for b, t, h, bt, _ in SOLVE_TRIL_CASES], +) +def test_costmodel_solve_tril(batch, sequence_length, heads, block, documented_us, tmp_path): chunks = sequence_length // block logical_programs = batch * chunks * heads torch.manual_seed(1) lower = torch.tril(torch.randn((block, block), dtype=torch.float32), diagonal=-1) * 0.01 - source = lower.reshape(1, block, 1, block).expand(chunks, block, heads, - block).reshape(batch, sequence_length, heads, block) - a = source.to(device="npu", dtype=torch.float32) - output = torch.zeros((batch, sequence_length, heads, block), dtype=torch.float32, device="npu") + a = torch.empty((batch, sequence_length, heads, block), dtype=torch.float32, device="npu") + a.view(batch, chunks, block, heads, block).copy_(lower.reshape(1, 1, block, 1, block).to("npu")) + output = torch.zeros_like(a) report_path = tmp_path / "solve_tril_route.json" def launch(): @@ -318,16 +333,30 @@ def launch(): ) launch() - inverse = torch.linalg.inv(torch.eye(block) + lower) - expected = inverse.reshape(1, block, 1, block).expand(chunks, block, heads, block).reshape_as(output).to("npu") - torch.testing.assert_close(output, expected, rtol=3e-2, atol=3e-2) + inverse = torch.linalg.inv(torch.eye(block) + lower).to("npu") + output_blocks = output.view(batch, chunks, block, heads, block) + probes = { + (0, 0, 0), + (batch // 2, chunks // 2, heads // 2), + (batch - 1, chunks - 1, heads - 1), + } + for batch_id, chunk_id, head_id in probes: + torch.testing.assert_close( + output_blocks[batch_id, chunk_id, :, head_id, :], + inverse, + rtol=3e-2, + atol=3e-2, + ) report = _load_route_report(report_path, "mixed_simd_simt") assert report["materialized_simt_anchor_count"] > 0 assert report["selected_superblock_factor"] == 4 assert report["effective_runtime_factor"] == 4 assert report["full_group_count"] == logical_programs // 4 assert report["tail_count"] == 0 - _assert_performance("solve_tril", launch, tmp_path / "solve_profile", 197.198) + case = f"solve_tril_B{batch}_T{sequence_length}_H{heads}_BT{block}" + _assert_performance(case, launch, tmp_path / "solve_profile", documented_us) + del output_blocks, inverse, output, a + torch.npu.empty_cache() @triton.jit From 9b9edc125a8b77b0b3f22c8a0e6e62ebcc8f917a Mon Sep 17 00:00:00 2001 From: yangkaixin Date: Tue, 1 Sep 2026 15:23:52 +0800 Subject: [PATCH 4/5] removed 'phase' concept from stage cost model --- third_party/ascend/backend/compiler.py | 3 + .../AscendModel/Analysis/StagePartitioner.h | 70 +- .../RouteModel/SimdSimtCostModel.h | 21 +- .../AscendModel/RouteModel/StageCostModels.h | 19 +- .../RouteModel/StageRouteCostModel.h | 23 +- .../AscendModel/Analysis/StagePartitioner.cpp | 1461 +++++++---------- .../RouteModel/SimdSimtCostModel.cpp | 207 +-- .../RouteModel/StageCostModels.cpp | 186 ++- .../RouteModel/StageRouteCostModel.cpp | 125 +- .../Transforms/SelectSimdSimtCostModel.cpp | 60 +- .../microbench/ascend_davidv100_v1.json | 12 +- .../data_provider/camodel/README.md | 123 +- .../camodel/compare_stage_costs.py | 425 +++++ .../camodel/parse_camodel_counts.py | 228 +++ .../simd_simt/david_v100_simd_simt_v1.json | 4 +- .../simd_simt/simd_simt_profile_schema.json | 2 +- .../General/TritonToLinalg/simt_scope.mlir | 14 +- .../unittest/costmodel_ut/PassesTest.cpp | 20 +- .../costmodel_ut/SimdSimtCostModelTest.cpp | 455 +++-- .../test_compiler_costmodel_contract.py | 14 + .../test_simd_simt_costmodel_cases.py | 24 +- 21 files changed, 2059 insertions(+), 1437 deletions(-) create mode 100644 third_party/ascend/costmodel/profiles/microbench/data_provider/camodel/compare_stage_costs.py create mode 100755 third_party/ascend/costmodel/profiles/microbench/data_provider/camodel/parse_camodel_counts.py diff --git a/third_party/ascend/backend/compiler.py b/third_party/ascend/backend/compiler.py index 309d3f5543..3c0faabdd4 100644 --- a/third_party/ascend/backend/compiler.py +++ b/third_party/ascend/backend/compiler.py @@ -1862,6 +1862,9 @@ def ttir_to_npubin(mod, metadata, opt): selected_factor = _selected_npuir_superblock_factor(metadata, opt) _compile_option_list += [f"--super-block-factor={selected_factor}"] + if not _is_debug_line_info_disabled(): + _compile_option_list += ["--enable-debug-info=true"] + npu_compiler_path, env = _get_npucompiler_path() cmd_list = ([npu_compiler_path, src_path] + _compile_option_list + ["-o", bin_file]) ret = subprocess.run(cmd_list, env=env, capture_output=True, check=True) diff --git a/third_party/ascend/costmodel/include/AscendModel/Analysis/StagePartitioner.h b/third_party/ascend/costmodel/include/AscendModel/Analysis/StagePartitioner.h index fb1d31b7ca..4e8f7cb568 100644 --- a/third_party/ascend/costmodel/include/AscendModel/Analysis/StagePartitioner.h +++ b/third_party/ascend/costmodel/include/AscendModel/Analysis/StagePartitioner.h @@ -1,4 +1,4 @@ -//===- StagePartitioner.h - Build semantic Phase/Stage IR ----*- C++ -*-===// +//===- StagePartitioner.h - Build semantic Stage IR ----------*- C++ -*-===// #ifndef ASCENDMODEL_ANALYSIS_STAGEPARTITIONER_H #define ASCENDMODEL_ANALYSIS_STAGEPARTITIONER_H @@ -9,7 +9,6 @@ #include "llvm/Support/Error.h" #include -#include #include #include @@ -21,40 +20,12 @@ struct StagePartitionerOptions { bool scopeSuperblockMaterializable = false; }; -enum class PhaseBoundaryDomain { - TriangularRecurrence, - LoadedIndexRowwiseReduction, - IndirectUnderfilledDot, -}; - -/// Result of PhaseBoundaryAnalysis. This is structural boundary evidence, -/// not a route or a cost-model decision. -struct PhaseBoundaryPlan { - PhaseBoundaryDomain domain; - std::string domainName; - std::optional triangularSolve; - /// Top-level semantic TTIR operations in execution order. Nested region - /// operations are owned transitively by their top-level root. - std::vector rootOperations; - /// Algorithm Phase ownership parallel to rootOperations. This is produced - /// by PhaseBoundaryAnalysis and is immutable input to StageBoundaryAnalysis; - /// Stage partitioning must not move a root across this boundary. - std::vector rootPhaseIds; - std::vector localSimtAnchorRoots; - - bool hasOperationGraph() const { - return !rootOperations.empty() && - rootOperations.size() == rootPhaseIds.size(); - } -}; - -/// Ordered post-transform TTIR roots plus the exact roots covered by each -/// materializable SIMT anchor. AutoBlockify V1's outer loop is represented -/// as a scheduling shell while its direct body operations remain semantic -/// roots; this prevents double ownership. +/// Ordered post-transform TTIR semantic roots. AutoBlockify V1's outer loop +/// is represented as a scheduling shell while its direct body operations +/// remain semantic roots; this prevents double ownership. Exact local SIMT +/// ownership is derived separately from the immutable SimtAnchorPlan. struct ProgramStructure { std::vector rootOperations; - std::vector localSimtAnchorRoots; }; class ProgramStructureAnalysis { @@ -63,26 +34,14 @@ class ProgramStructureAnalysis { analyze(ModuleOp module, const SimtAnchorPlan &anchorPlan) const; }; -/// Recognizes algorithm-level serial regions. The current feature-summary -/// overload is an explicit fallback; the operation-graph overload is the -/// target implementation for production Stage ownership. -class PhaseBoundaryAnalysis { -public: - llvm::Expected> - analyze(ModuleOp module, const SimtAnchorPlan &anchorPlan, - const SimdSimtFeatureSummary &features, - const StagePartitionerOptions &options) const; -}; - -/// Splits each Phase into single-kind Stages. It does not evaluate cycles or -/// choose SIMD/SIMT. The optional anchor plan is used only as exact ownership -/// evidence for materializable local SIMT Stages. +/// Splits ordered semantic roots directly into single-kind Stages. It does +/// not evaluate cycles or choose SIMD/SIMT. The anchor plan is used only as +/// exact ownership evidence for materializable local SIMT Stages. class StageBoundaryAnalysis { public: llvm::Expected - analyze(const PhaseBoundaryPlan &phasePlan, - const SimdSimtFeatureSummary &features, - const SimtAnchorPlan *anchorPlan = nullptr) const; + analyze(const ProgramStructure &structure, + const SimtAnchorPlan &anchorPlan) const; }; /// Derives structural facts for every already-owned Stage. It never chooses @@ -122,14 +81,13 @@ class StagePartitionVerifier { llvm::Error verify(const StagePartition &partition) const; }; -/// Partitions post-layout/post-AutoBlockify-V1 TTIR facts into serial Phases -/// and single-mode Stages. A Stage may later be evaluated as SIMD or SIMT, -/// but it is never internally mixed. +/// Partitions post-layout/post-AutoBlockify-V1 TTIR directly into ordered, +/// single-mode Stages. A Stage may later be evaluated as SIMD or SIMT, but it +/// is never internally mixed. class StagePartitioner { public: - llvm::Expected> + llvm::Expected partition(ModuleOp module, const SimtAnchorPlan &anchorPlan, - const SimdSimtFeatureSummary &features, const StagePartitionerOptions &options) const; }; diff --git a/third_party/ascend/costmodel/include/AscendModel/RouteModel/SimdSimtCostModel.h b/third_party/ascend/costmodel/include/AscendModel/RouteModel/SimdSimtCostModel.h index 6991cabf28..0122367b89 100644 --- a/third_party/ascend/costmodel/include/AscendModel/RouteModel/SimdSimtCostModel.h +++ b/third_party/ascend/costmodel/include/AscendModel/RouteModel/SimdSimtCostModel.h @@ -35,35 +35,20 @@ llvm::StringRef stringifySimdSimtCandidate(SimdSimtCandidateKind candidate); /// contains no Operation pointers, so reports remain stable and serializable. struct SimtAnchorFeatureSummary { int64_t count = 0; - int64_t conditionalBranchCount = 0; - int64_t divergentBranchCount = 0; - double activeLaneRatio = 1.0; std::vector triangularSolves; CandidateLowerability kernelLowerability; llvm::json::Object toJSON() const; }; -/// Static, workload-name-independent properties extracted directly from a -/// generic TTIR ModuleOp. Weighted fields include statically known scf.for -/// trip counts. +/// Kernel-level facts needed before Stage partitioning. Operation workload +/// belongs to StageFeatureAnalysis/StageWorkloadAnalysis and is intentionally +/// not duplicated here. struct SimdSimtFeatureSummary { - int64_t loadOps = 0; - int64_t storeOps = 0; - int64_t reduceOps = 0; - int64_t dotOps = 0; - int64_t loadedIndexDependentMemoryOps = 0; - int64_t dotFlops = 0; - int64_t staticLoopTripCountMax = 1; - int64_t conditionalBranchCount = 0; - int64_t divergentBranchCount = 0; - double activeLaneRatio = 1.0; - /// Scheduling/layout facts read from the transformed TTIR consumed by this /// model. These make it explicit that layout merging and AutoBlockify V1 /// ran before feature extraction rather than being guessed from source TTIR. bool autoBlockifyV1Applied = false; - int64_t autoBlockifyV1LoopCount = 0; bool hasExplicitScope = false; SimtAnchorFeatureSummary simtAnchors; diff --git a/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageCostModels.h b/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageCostModels.h index c1f3aaa097..aa9db382d6 100644 --- a/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageCostModels.h +++ b/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageCostModels.h @@ -89,20 +89,18 @@ struct LogicalStage { /// True when this Stage has exact operation ownership/live-in/live-out and /// can therefore become a local SIMT scope inside a mixed kernel. bool localSimtMaterializable = false; + /// True when the one selected local scope will be a direct operation of an + /// AutoBlockify V1 loop body. NPUIR's current scope-SuperBlock ABI requires + /// this stronger condition for F2/F4; nested scopes remain legal at F1. + bool localSuperblockMaterializable = false; std::vector legalSimtFactors; std::vector localSimtFactors; }; -struct LogicalPhase { - std::string id; - std::vector stages; -}; - struct StagePartition { - std::string domain; bool operationOwnershipComplete = false; int64_t modeledOperationCount = 0; - std::vector phases; + std::vector stages; }; struct StageOperationRate { @@ -144,9 +142,10 @@ struct HardwareProfile { /// option, not a hardware constant, and bounds cross-group interleaving in /// recurrence Stage models. int64_t logicalWarpGroupCount = 1; - /// SuperBlock hides latency, but large factors also replicate long-lived - /// recurrence state. These target-profile values model the resulting - /// register/stack pressure without naming a workload. + /// Long-lived recurrence state consumes finite register/stack bandwidth. + /// The byte rate is shared by the SIMD recurrence-state term and the extra + /// pressure created when a SIMT SuperBlock replicates that state; neither + /// formula depends on a workload name. /// Largest factor that still gives proportional latency-hiding benefit. int64_t superblockUsefulFactorLimit = 1; /// Largest factor that may replicate loop-carried live state without an diff --git a/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageRouteCostModel.h b/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageRouteCostModel.h index 92a3dab1f2..9de0c43732 100644 --- a/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageRouteCostModel.h +++ b/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageRouteCostModel.h @@ -57,6 +57,11 @@ struct StageModelFeatures { bool hasReduction = false; bool hasDot = false; bool hasConversionPack = false; + /// True when the Stage is part of the logical-program body created by + /// AutoBlockify V1. A factor-F local SuperBlock executes SIMD Stages in + /// this body once for every grouped logical program, while the selected + /// local SIMT Stage consumes the factor through its SuperBlock model. + bool replicatedByLocalSuperBlock = false; int64_t conditionalBranchCount = 0; int64_t divergentBranchCount = 0; int64_t loopBackedgeCount = 0; @@ -134,6 +139,11 @@ struct LogicalStageCost { StageModelFeatures features; StageWorkload workload; int64_t ownedOperationCount = 0; + /// Unique source locations of the TTIR operations owned by this Stage. + /// These are calibration provenance only: they let a debug-line-enabled + /// CaModel artifact map binary PCs back to the immutable StagePartition + /// without adding marker operations or attributes to production IR. + std::vector sourceLocations; int64_t liveInCount = 0; int64_t liveOutCount = 0; /// Static tensor footprint crossing the Stage boundary. Counts alone are @@ -149,6 +159,7 @@ struct LogicalStageCost { int64_t scopeOutputTensorBytes = 0; std::vector simtAnchorIndices; bool localSimtMaterializable = false; + bool localSuperblockMaterializable = false; /// Factors legal for a whole-kernel pure-SIMT schedule. std::vector legalSimtFactors; /// Factors legal when this Stage alone is materialized as a local scope. @@ -158,21 +169,12 @@ struct LogicalStageCost { llvm::json::Object toJSON() const; }; -struct LogicalPhaseCost { - std::string id; - std::vector stages; - - llvm::json::Object toJSON() const; -}; - struct StageCostTable { - std::string domain; bool operationOwnershipComplete = false; int64_t modeledOperationCount = 0; std::string profileVersion; int64_t logicalProgramCountHint = 0; int64_t physicalCoreCountHint = 0; - std::vector phases; std::vector stages; }; @@ -199,7 +201,6 @@ struct StageRoutePlan { std::vector implementations; std::vector entryTransitionCycles; std::vector logicalStageCycles; - std::vector logicalPhaseCycles; int64_t routeSuperblockFactor = 1; int64_t runtimePhysicalProgramCount = 0; int64_t runtimeWaveCount = 1; @@ -210,11 +211,9 @@ struct StageRoutePlan { struct StageCostModelSummary { bool applied = false; - std::string domain; bool operationOwnershipComplete = false; int64_t modeledOperationCount = 0; std::string profileVersion; - std::vector phases; std::vector stages; StageTransitionCost transition; StageRoutePlan allSimd; diff --git a/third_party/ascend/costmodel/lib/AscendModel/Analysis/StagePartitioner.cpp b/third_party/ascend/costmodel/lib/AscendModel/Analysis/StagePartitioner.cpp index bc2b13fec7..c30acd77a4 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/Analysis/StagePartitioner.cpp +++ b/third_party/ascend/costmodel/lib/AscendModel/Analysis/StagePartitioner.cpp @@ -1,4 +1,4 @@ -//===- StagePartitioner.cpp - Build semantic Phase/Stage IR -------------===// +//===- StagePartitioner.cpp - Build semantic Stage IR -------------------===// #include "AscendModel/Analysis/StagePartitioner.h" @@ -8,11 +8,13 @@ #include "llvm/ADT/SetVector.h" #include "llvm/ADT/StringSet.h" #include "llvm/ADT/StringSwitch.h" +#include "llvm/ADT/Twine.h" #include "llvm/Support/raw_ostream.h" #include #include #include +#include #include #include @@ -331,214 +333,14 @@ static void mergeWorkload(StageWorkload &into, StageWorkload from) { static void makePerIteration(LogicalStage &stage) { const double count = static_cast(std::max(1, stage.iterationCount)); - StageWorkload &work = stage.workload; - work.scalarOperations /= count; - work.loadBytes /= count; - work.storeBytes /= count; - work.loadWarpInstructions /= count; - work.storeWarpInstructions /= count; - work.predicateElements /= count; - work.shuffleLaneSteps /= count; - work.dotFlops /= count; - work.estimatedSpillTransactions /= count; - for (auto &entry : work.operationElements) - entry.second /= count; - recomputeIssueElements(work); -} - -static LogicalStage makeStage(llvm::StringRef id, StageCostModelKind kind, - StageScheduleKind schedule, int64_t iterations, - StageWorkload workload) { - LogicalStage stage; - stage.id = id.str(); - stage.costModelKind = kind; - stage.scheduleKind = schedule; - stage.iterationCount = std::max(1, iterations); - stage.workload = std::move(workload); - makePerIteration(stage); - return stage; -} - -static LogicalStage withControl(LogicalStage stage, int64_t branches, - int64_t divergent, double activeLaneRatio) { - stage.features.conditionalBranchCount = std::max(0, branches); - stage.features.divergentBranchCount = std::max(0, divergent); - stage.features.activeLaneRatio = std::clamp(activeLaneRatio, 0.0, 1.0); - return stage; -} - -static LogicalStage asLocalSIMT(LogicalStage stage) { - stage.localSimtMaterializable = true; - // StageModeLegalityAnalysis opens F2/F4 only when backend integration has - // the AutoBlockify V1 Scope SuperBlock wrapper available. - stage.localSimtFactors = {1}; - return stage; -} - -static void addPhase(StagePartition &partition, llvm::StringRef id, - LogicalStage stage) { - LogicalPhase phase; - phase.id = id.str(); - phase.stages.push_back(std::move(stage)); - partition.phases.push_back(std::move(phase)); -} - -static bool operationTreeContainsName(Operation *root, llvm::StringRef name); - -static bool hasPhase(const PhaseBoundaryPlan *plan, llvm::StringRef id) { - return plan && llvm::is_contained(plan->rootPhaseIds, id); -} - -static void prependAutoBlockifyStages(StagePartition &partition, - const SimdSimtFeatureSummary &features, - const PhaseBoundaryPlan *plan) { - if (!features.autoBlockifyV1Applied && - !hasPhase(plan, "auto_blockify_dispatch")) - return; - LogicalPhase phase; - phase.id = "auto_blockify_dispatch"; - phase.stages.push_back(makeStage("physical_program_dispatch", - StageCostModelKind::AutoBlockifyDispatch, - StageScheduleKind::StraightLine, 1, {})); - if (features.autoBlockifyV1LoopCount > 0 || - hasPhase(plan, "auto_blockify_dispatch")) - phase.stages.push_back( - makeStage("logical_program_loop", StageCostModelKind::AutoBlockifyLoop, - StageScheduleKind::IndependentPipelined, - std::max(1, features.autoBlockifyV1LoopCount), {})); - partition.phases.push_back(std::move(phase)); -} - -static StagePartition -partitionTriangular(const SimdSimtFeatureSummary &features, - const TriangularSolveFacts &facts, - const PhaseBoundaryPlan *plan) { - StagePartition partition; - partition.domain = "triangular_recurrence"; - prependAutoBlockifyStages(partition, features, plan); - - LogicalStage head = withControl( - makeStage("head_index_mask", StageCostModelKind::PredicateMask, - StageScheduleKind::StraightLine, 1, {}), - features.conditionalBranchCount - - features.simtAnchors.conditionalBranchCount, - features.divergentBranchCount - features.simtAnchors.divergentBranchCount, - features.activeLaneRatio); - head.workload.paysKernelSetup = true; - addPhase(partition, "head", std::move(head)); - - if (hasPhase(plan, "diagonal_load")) - addPhase(partition, "diagonal_load", - makeStage("load_diagonal_tiles", - StageCostModelKind::ContinuousTileMemory, - StageScheduleKind::IndependentPipelined, 1, {})); - - const int64_t iterations = std::max(1, facts.recurrenceLoopCount); - LogicalStage recurrence = asLocalSIMT(withControl( - makeStage("diagonal_inverse_recurrence", - StageCostModelKind::LoopCarriedRecurrence, - StageScheduleKind::LoopCarriedSerial, iterations, {}), - features.simtAnchors.conditionalBranchCount, - features.simtAnchors.divergentBranchCount, - features.simtAnchors.activeLaneRatio)); - const int64_t rows = - std::max(1, facts.blockRows - facts.recurrenceStartRow); - recurrence.features.parallelRecurrenceGroupCount = - std::max(1, (iterations + rows - 1) / rows); - addPhase(partition, "diagonal_inverse", std::move(recurrence)); - - if (hasPhase(plan, "merge_store")) { - bool hasDot = false; - bool hasStore = false; - for (auto root : llvm::enumerate(plan->rootOperations)) { - if (plan->rootPhaseIds[root.index()] != "merge_store") - continue; - hasDot |= operationTreeContainsName(root.value(), "tt.dot"); - hasStore |= operationTreeContainsName(root.value(), "tt.store"); - } - LogicalPhase phase; - phase.id = "merge_store"; - if (hasDot) - phase.stages.push_back( - makeStage("dense_dot_tail", StageCostModelKind::CubeRoofline, - StageScheduleKind::IndependentPipelined, - std::max(1, facts.denseDotTailOps), {})); - if (hasStore) - phase.stages.push_back(makeStage("store_inverse_tile", - StageCostModelKind::ContinuousTileStore, - StageScheduleKind::StraightLine, 1, {})); - if (!phase.stages.empty()) - partition.phases.push_back(std::move(phase)); - } - return partition; -} - -static StagePartition partitionRowwise(const SimdSimtFeatureSummary &features, - const PhaseBoundaryPlan *plan) { - StagePartition partition; - partition.domain = "loaded_index_rowwise_reduction"; - prependAutoBlockifyStages(partition, features, plan); - addPhase(partition, "row_dispatch", - makeStage("row_index_generation", - StageCostModelKind::IndexGeneration, - StageScheduleKind::StraightLine, 1, {})); - addPhase(partition, "row_load", - asLocalSIMT(makeStage( - "indirect_row_gather", StageCostModelKind::IndirectGatherMemory, - StageScheduleKind::PartiallyDependent, 1, {}))); - const int64_t iterations = - std::max(1, features.staticLoopTripCountMax); - addPhase(partition, "row_reduction", - makeStage("rowwise_reduction", StageCostModelKind::RowwiseReduction, - StageScheduleKind::PartiallyDependent, iterations, {})); - addPhase(partition, "convert_store", - makeStage("conversion_pack_store", - StageCostModelKind::ConversionPack, - StageScheduleKind::IndependentPipelined, iterations, {})); - return partition; -} - -static StagePartition -partitionIndirectDot(const SimdSimtFeatureSummary &features, - const PhaseBoundaryPlan *plan) { - StagePartition partition; - partition.domain = "indirect_underfilled_dot"; - prependAutoBlockifyStages(partition, features, plan); - addPhase(partition, "index_setup", - makeStage("index_generation", StageCostModelKind::IndexGeneration, - StageScheduleKind::StraightLine, 1, {})); - addPhase(partition, "gather_tiles", - asLocalSIMT(makeStage( - "indirect_tile_gather", StageCostModelKind::IndirectGatherMemory, - StageScheduleKind::PartiallyDependent, 1, {}))); - addPhase(partition, "dot", - makeStage("tiny_cube_dot", StageCostModelKind::TinyCubeRoofline, - StageScheduleKind::IndependentPipelined, 1, {})); - addPhase(partition, "output_store", - makeStage("store_dot_result", - StageCostModelKind::ContinuousTileStore, - StageScheduleKind::StraightLine, 1, {})); - return partition; -} - -static bool anchorMatchesStage(const SimtAnchorDescriptor &anchor, - const LogicalStage &stage) { - if (!anchor.materializable || !stage.localSimtMaterializable) - return false; - if (stage.costModelKind == StageCostModelKind::LoopCarriedRecurrence) - return anchor.kind == SimtAnchorKind::TriangularSolveLoop; - if (stage.costModelKind == StageCostModelKind::IndirectGatherMemory || - stage.costModelKind == StageCostModelKind::IndirectScalarMemory) - return anchor.kind == SimtAnchorKind::DirectGather || - anchor.kind == SimtAnchorKind::LoadedIndexDependentMemory; - return false; + scaleWorkload(stage.workload, 1.0 / count); } static Operation *getTopLevelSemanticRoot(Operation *operation); static bool stageOwnsAnchor(const LogicalStage &stage, const SimtAnchorDescriptor &anchor) { - if (!anchorMatchesStage(anchor, stage)) + if (!anchor.materializable || !stage.localSimtMaterializable) return false; auto owns = [&](Operation *operation) { Operation *root = getTopLevelSemanticRoot(operation); @@ -555,21 +357,32 @@ static bool stageOwnsAnchor(const LogicalStage &stage, /// second source of Stage boundaries. static void attachExactAnchorOwnership(StagePartition &partition, const SimtAnchorPlan &anchorPlan) { - for (LogicalPhase &phase : partition.phases) { - for (LogicalStage &stage : phase.stages) { - if (!stage.localSimtMaterializable) - continue; - stage.simtAnchorIndices.clear(); - for (auto indexedAnchor : llvm::enumerate(anchorPlan.anchors)) { - const SimtAnchorDescriptor &anchor = indexedAnchor.value(); - if (anchor.materializable && stageOwnsAnchor(stage, anchor)) - stage.simtAnchorIndices.push_back( - static_cast(indexedAnchor.index())); + for (LogicalStage &stage : partition.stages) { + if (!stage.localSimtMaterializable) + continue; + stage.simtAnchorIndices.clear(); + stage.localSuperblockMaterializable = false; + bool allAnchorsDirectlyOwnedByV1Loop = true; + for (auto indexedAnchor : llvm::enumerate(anchorPlan.anchors)) { + const SimtAnchorDescriptor &anchor = indexedAnchor.value(); + if (anchor.materializable && stageOwnsAnchor(stage, anchor)) { + stage.simtAnchorIndices.push_back( + static_cast(indexedAnchor.index())); + Operation *insertionPoint = anchor.scopeOperations.size() > 1 + ? anchor.scopeInsertionPoint + : anchor.operation; + Operation *blockOwner = insertionPoint && insertionPoint->getBlock() + ? insertionPoint->getBlock()->getParentOp() + : nullptr; + allAnchorsDirectlyOwnedByV1Loop &= + blockOwner && blockOwner->hasAttr("ta.auto_blockify_v1.loop"); } - stage.localSimtMaterializable = !stage.simtAnchorIndices.empty(); - if (!stage.localSimtMaterializable) - stage.localSimtFactors.clear(); } + stage.localSimtMaterializable = !stage.simtAnchorIndices.empty(); + stage.localSuperblockMaterializable = + stage.simtAnchorIndices.size() == 1 && allAnchorsDirectlyOwnedByV1Loop; + if (!stage.localSimtMaterializable) + stage.localSimtFactors.clear(); } } @@ -593,6 +406,16 @@ static Operation *getTopLevelSemanticRoot(Operation *operation) { return nullptr; } +static bool isInsideAutoBlockifyV1Loop(Operation *operation) { + if (!operation || operation->hasAttr("ta.auto_blockify_v1.loop")) + return false; + for (Operation *parent = operation->getParentOp(); parent; + parent = parent->getParentOp()) + if (parent->hasAttr("ta.auto_blockify_v1.loop")) + return true; + return false; +} + static std::vector collectTopLevelSemanticRoots(ModuleOp module) { std::vector result; auto appendBlock = [&](Block &block) { @@ -643,235 +466,118 @@ static bool operationTreeContainsLoadedIndexMemory(Operation *root) { return found; } -/// PhaseBoundaryAnalysis owns the algorithm-level serial cut. Each root is -/// assigned exactly one Phase id in execution order. The state machines are -/// monotone: after a boundary is crossed, a later root cannot move back to an -/// earlier Phase. Cost and candidate mode are intentionally absent here. -static llvm::Error assignRootPhaseIds(PhaseBoundaryPlan &plan) { - llvm::DenseSet anchorRoots(plan.localSimtAnchorRoots.begin(), - plan.localSimtAnchorRoots.end()); - std::optional firstAnchorIndex; - std::optional lastAnchorIndex; - for (auto indexedRoot : llvm::enumerate(plan.rootOperations)) { - if (!anchorRoots.contains(indexedRoot.value())) - continue; - if (!firstAnchorIndex) - firstAnchorIndex = indexedRoot.index(); - lastAnchorIndex = indexedRoot.index(); - } - if (plan.domain == PhaseBoundaryDomain::TriangularRecurrence && - !firstAnchorIndex) - return llvm::createStringError( - std::errc::invalid_argument, - "triangular PhaseBoundaryAnalysis requires an exact anchor region"); - if (plan.domain == PhaseBoundaryDomain::TriangularRecurrence && - firstAnchorIndex && lastAnchorIndex) { - for (size_t index = *firstAnchorIndex; index <= *lastAnchorIndex; ++index) - if (!anchorRoots.contains(plan.rootOperations[index])) - return llvm::createStringError( - std::errc::invalid_argument, - "triangular PhaseBoundaryAnalysis requires a contiguous planned " - "scope region"); - } - - plan.rootPhaseIds.clear(); - plan.rootPhaseIds.reserve(plan.rootOperations.size()); - llvm::StringRef current; - for (auto indexedRoot : llvm::enumerate(plan.rootOperations)) { - Operation *root = indexedRoot.value(); - if (!root) - return llvm::createStringError( - std::errc::invalid_argument, - "PhaseBoundaryAnalysis received a null semantic root"); - if (root->hasAttr("ta.auto_blockify_v1.loop") || - root->hasAttr("ta.auto_blockify_v1.schedule")) { - plan.rootPhaseIds.push_back("auto_blockify_dispatch"); - continue; - } - - switch (plan.domain) { - case PhaseBoundaryDomain::TriangularRecurrence: - if (firstAnchorIndex && indexedRoot.index() >= *firstAnchorIndex && - indexedRoot.index() <= *lastAnchorIndex) - current = "diagonal_inverse"; - else if (lastAnchorIndex && indexedRoot.index() > *lastAnchorIndex) - current = "merge_store"; - else if (operationTreeContainsName(root, "tt.dot") || - operationTreeContainsName(root, "tt.store")) - current = "merge_store"; - else if ((current.empty() || current == "head") && - operationTreeContainsName(root, "tt.load")) - current = "diagonal_load"; - else if (current.empty()) - current = "head"; - break; - case PhaseBoundaryDomain::LoadedIndexRowwiseReduction: - if (operationTreeContainsName(root, "tt.reduce")) - current = "row_reduction"; - else if (current == "row_reduction" || - operationTreeContainsName(root, "tt.store")) - current = "convert_store"; - else if ((current.empty() || current == "row_dispatch") && - operationTreeContainsLoadedIndexMemory(root)) - current = "row_load"; - else if (current.empty()) - current = "row_dispatch"; - break; - case PhaseBoundaryDomain::IndirectUnderfilledDot: - if (operationTreeContainsName(root, "tt.dot")) - current = "dot"; - else if (current == "dot" || operationTreeContainsName(root, "tt.store")) - current = "output_store"; - else if ((current.empty() || current == "index_setup") && - operationTreeContainsLoadedIndexMemory(root)) - current = "gather_tiles"; - else if (current.empty()) - current = "index_setup"; - break; +static bool operationTreeHasTrueLoopCarriedDependency(Operation *root) { + bool found = false; + if (!root) + return found; + root->walk([&](Operation *operation) { + if (found || operation->hasAttr("ta.auto_blockify_v1.loop")) + return; + const llvm::StringRef name = operation->getName().getStringRef(); + if ((name != "scf.for" && name != "scf.while") || + operation->getNumRegions() == 0 || operation->getRegion(0).empty()) + return; + Block &body = operation->getRegion(0).front(); + const unsigned firstCarriedArgument = name == "scf.for" ? 1 : 0; + for (unsigned index = firstCarriedArgument; index < body.getNumArguments(); + ++index) { + BlockArgument argument = body.getArgument(index); + if (!argument.use_empty() && !isPointerLikeType(argument.getType()) && + !isAddressOnlyLoopValue(argument)) { + found = true; + return; + } } - plan.rootPhaseIds.push_back(current.str()); - } - if (plan.rootPhaseIds.size() != plan.rootOperations.size()) - return llvm::createStringError( - std::errc::invalid_argument, - "PhaseBoundaryAnalysis did not own every semantic root"); - llvm::StringSet<> closedPhases; - llvm::StringRef currentPhase; - for (const std::string &phaseId : plan.rootPhaseIds) { - if (phaseId == currentPhase) - continue; - if (!currentPhase.empty()) - closedPhases.insert(currentPhase); - if (closedPhases.contains(phaseId)) - return llvm::createStringError( - std::errc::invalid_argument, - "PhaseBoundaryAnalysis produced a non-contiguous Phase '%s'", - phaseId.c_str()); - currentPhase = phaseId; - } - return llvm::Error::success(); + }); + return found; } -static LogicalStage *findStage(StagePartition &partition, llvm::StringRef id) { - for (LogicalPhase &phase : partition.phases) - for (LogicalStage &stage : phase.stages) - if (stage.id == id) - return &stage; - return nullptr; +static bool operationTreeHasAnyName(Operation *root, + llvm::ArrayRef names) { + return llvm::any_of(names, [&](llvm::StringRef name) { + return operationTreeContainsName(root, name); + }); } -static llvm::StringRef stageIdForPhase(PhaseBoundaryDomain domain, - llvm::StringRef phaseId) { - if (domain == PhaseBoundaryDomain::TriangularRecurrence) - return llvm::StringSwitch(phaseId) - .Case("head", "head_index_mask") - .Case("diagonal_load", "load_diagonal_tiles") - .Case("diagonal_inverse", "diagonal_inverse_recurrence") - .Default({}); - if (domain == PhaseBoundaryDomain::LoadedIndexRowwiseReduction) - return llvm::StringSwitch(phaseId) - .Case("row_dispatch", "row_index_generation") - .Case("row_load", "indirect_row_gather") - .Case("row_reduction", "rowwise_reduction") - .Case("convert_store", "conversion_pack_store") - .Default({}); - return llvm::StringSwitch(phaseId) - .Case("index_setup", "index_generation") - .Case("gather_tiles", "indirect_tile_gather") - .Case("dot", "tiny_cube_dot") - .Case("output_store", "store_dot_result") - .Default({}); +/// Classify one transitive semantic ownership unit. This function consumes +/// only TTIR structure; it does not inspect a kernel name, workload name, +/// measured performance, or route score. +static StageCostModelKind classifySemanticRoot(Operation *root) { + if (root->hasAttr("ta.auto_blockify_v1.loop")) + return StageCostModelKind::AutoBlockifyLoop; + if (root->hasAttr("ta.auto_blockify_v1.schedule")) + return StageCostModelKind::AutoBlockifyDispatch; + if (operationTreeHasTrueLoopCarriedDependency(root)) + return StageCostModelKind::LoopCarriedRecurrence; + if (operationTreeHasAnyName(root, {"tt.reduce", "tt.scan", "linalg.reduce"})) + return StageCostModelKind::RowwiseReduction; + if (operationTreeHasAnyName(root, {"tt.dot"})) + return StageCostModelKind::CubeRoofline; + if (operationTreeContainsLoadedIndexMemory(root) || + operationTreeHasAnyName(root, {"tt.gather"})) + return StageCostModelKind::IndirectGatherMemory; + if (operationTreeHasAnyName(root, {"scf.for", "scf.while"})) + return StageCostModelKind::IndependentPipelinedLoop; + if (operationTreeHasAnyName( + root, {"tt.fp_to_fp", "arith.extf", "arith.truncf", "arith.fptosi", + "arith.fptoui", "arith.sitofp", "arith.uitofp"})) + return StageCostModelKind::ConversionPack; + const bool hasLoad = operationTreeHasAnyName(root, {"tt.load"}); + const bool hasStore = operationTreeHasAnyName(root, {"tt.store"}); + if (hasStore && !hasLoad) + return StageCostModelKind::ContinuousTileStore; + if (hasLoad || hasStore) + return StageCostModelKind::ContinuousTileMemory; + if (operationTreeHasAnyName(root, {"arith.cmpi", "arith.cmpf"})) + return StageCostModelKind::PredicateMask; + if (operationTreeHasAnyName(root, {"tt.get_program_id", "tt.addptr", + "tt.advance", "arith.index_cast"})) + return StageCostModelKind::IndexGeneration; + if (operationTreeHasAnyName(root, {"scf.if", "cf.cond_br"})) + return StageCostModelKind::ScalarControl; + return StageCostModelKind::ScalarIssue; } -static int64_t getStageOrdinal(const StagePartition &partition, - const LogicalStage *target) { - int64_t ordinal = 0; - for (const LogicalPhase &phase : partition.phases) { - for (const LogicalStage &stage : phase.stages) { - if (&stage == target) - return ordinal; - ++ordinal; - } - } - return -1; +static StageScheduleKind scheduleForSemanticRoot(Operation *root, + StageCostModelKind kind) { + if (kind == StageCostModelKind::LoopCarriedRecurrence) + return StageScheduleKind::LoopCarriedSerial; + if (kind == StageCostModelKind::IndirectGatherMemory) + return StageScheduleKind::PartiallyDependent; + if (kind == StageCostModelKind::AutoBlockifyLoop || + kind == StageCostModelKind::IndependentPipelinedLoop || + operationTreeHasAnyName(root, {"scf.for", "scf.while"})) + return StageScheduleKind::IndependentPipelined; + return StageScheduleKind::StraightLine; } -static llvm::Error assignOperation(LogicalStage *stage, Operation *operation, - llvm::DenseSet &owned) { - if (!stage || !operation) - return llvm::createStringError( - std::errc::invalid_argument, - "StageBoundaryAnalysis could not map a TTIR operation"); - if (!owned.insert(operation).second) - return llvm::createStringError( - std::errc::invalid_argument, - "StageBoundaryAnalysis assigned a TTIR operation more than once"); - stage->operations.push_back(operation); - return llvm::Error::success(); +static int64_t semanticRootIterationCount(Operation *root) { + int64_t iterations = 1; + if (!root || root->hasAttr("ta.auto_blockify_v1.loop")) + return iterations; + root->walk([&](Operation *operation) { + if (!operation->hasAttr("ta.auto_blockify_v1.loop")) + iterations = std::max(iterations, getLoopTripCount(operation, 1)); + }); + return iterations; } -static llvm::Error -attachCompleteOperationOwnership(StagePartition &partition, - const PhaseBoundaryPlan &plan) { - if (!plan.hasOperationGraph()) - return llvm::createStringError( - std::errc::invalid_argument, - "StageBoundaryAnalysis requires complete Phase root ownership"); - llvm::DenseSet owned; - int64_t lastStageOrdinal = -1; - bool mergeStoreReached = findStage(partition, "dense_dot_tail") == nullptr; - - for (auto indexedRoot : llvm::enumerate(plan.rootOperations)) { - Operation *root = indexedRoot.value(); - const llvm::StringRef phaseId = plan.rootPhaseIds[indexedRoot.index()]; - LogicalStage *target = nullptr; - if (phaseId == "auto_blockify_dispatch") { - if (root->hasAttr("ta.auto_blockify_v1.loop")) - target = findStage(partition, "logical_program_loop"); - else if (!root->hasAttr("ta.auto_blockify_v1.schedule")) - return llvm::createStringError( - std::errc::invalid_argument, - "AutoBlockify Phase contains a root without V1 provenance"); - Operation *parent = root->getParentOp(); - if (!target) - target = findStage(partition, - parent && parent->hasAttr("ta.auto_blockify_v1.loop") - ? "logical_program_loop" - : "physical_program_dispatch"); - } - - if (!target && plan.domain == PhaseBoundaryDomain::TriangularRecurrence && - phaseId == "merge_store") { - mergeStoreReached |= operationTreeContainsName(root, "tt.store"); - target = findStage(partition, mergeStoreReached ? "store_inverse_tile" - : "dense_dot_tail"); - } - if (!target) - target = findStage(partition, stageIdForPhase(plan.domain, phaseId)); - const int64_t ordinal = getStageOrdinal(partition, target); - if (ordinal < 0) - return llvm::createStringError( - std::errc::invalid_argument, - "StageBoundaryAnalysis selected a missing Stage for Phase '%s' " - "and root '%s'", - phaseId.str().c_str(), root->getName().getStringRef().str().c_str()); - if (ordinal < lastStageOrdinal) - return llvm::createStringError( - std::errc::invalid_argument, - "StageBoundaryAnalysis produced non-contiguous Stage ownership"); - lastStageOrdinal = ordinal; - if (llvm::Error error = assignOperation(target, root, owned)) - return error; - } +static bool hasOrderedStageBoundary(Operation *root) { + bool ordered = false; + if (!root) + return ordered; + root->walk([&](Operation *operation) { + const llvm::StringRef name = operation->getName().getStringRef(); + ordered |= operation->getNumRegions() > 0 || + name.starts_with("tt.atomic") || name.contains("barrier") || + name.contains("sync") || name == "cf.cond_br"; + }); + return ordered; +} - if (owned.size() != plan.rootOperations.size()) - return llvm::createStringError( - std::errc::invalid_argument, - "StageBoundaryAnalysis did not conserve TTIR operation ownership"); - partition.operationOwnershipComplete = true; - partition.modeledOperationCount = - static_cast(plan.rootOperations.size()); - return llvm::Error::success(); +static std::string makeStageId(size_t ordinal, StageCostModelKind kind) { + return ("stage_" + llvm::Twine(ordinal) + "_" + stringifyStageCostModel(kind)) + .str(); } static void collectOwnedOperationTree(Operation *root, @@ -931,29 +637,27 @@ static int64_t staticTensorBytes(llvm::ArrayRef values) { /// Values defined outside and consumed inside are live-ins; values defined /// inside and consumed by any operation outside are live-outs. static void deriveStageLiveValues(StagePartition &partition) { - for (LogicalPhase &phase : partition.phases) { - for (LogicalStage &stage : phase.stages) { - llvm::DenseSet owned; - for (Operation *root : stage.operations) - collectOwnedOperationTree(root, owned); - - llvm::SetVector liveIns; - llvm::SetVector liveOuts; - for (Operation *operation : owned) { - for (Value operand : operation->getOperands()) - if (!isValueDefinedInside(operand, owned)) - liveIns.insert(operand); - for (Value result : operation->getResults()) - if (llvm::any_of(result.getUsers(), [&](Operation *user) { - return !owned.contains(user); - })) - liveOuts.insert(result); - } - stage.liveIns.assign(liveIns.begin(), liveIns.end()); - stage.liveOuts.assign(liveOuts.begin(), liveOuts.end()); - stage.liveInBytes = staticTensorBytes(stage.liveIns); - stage.liveOutBytes = staticTensorBytes(stage.liveOuts); + for (LogicalStage &stage : partition.stages) { + llvm::DenseSet owned; + for (Operation *root : stage.operations) + collectOwnedOperationTree(root, owned); + + llvm::SetVector liveIns; + llvm::SetVector liveOuts; + for (Operation *operation : owned) { + for (Value operand : operation->getOperands()) + if (!isValueDefinedInside(operand, owned)) + liveIns.insert(operand); + for (Value result : operation->getResults()) + if (llvm::any_of(result.getUsers(), [&](Operation *user) { + return !owned.contains(user); + })) + liveOuts.insert(result); } + stage.liveIns.assign(liveIns.begin(), liveIns.end()); + stage.liveOuts.assign(liveOuts.begin(), liveOuts.end()); + stage.liveInBytes = staticTensorBytes(stage.liveIns); + stage.liveOutBytes = staticTensorBytes(stage.liveOuts); } } @@ -963,73 +667,68 @@ static void deriveStageLiveValues(StagePartition &partition) { /// charging its complete live-out footprint would invent UB traffic. static void deriveLocalSimtScopeTraffic(StagePartition &partition, const SimtAnchorPlan &anchorPlan) { - for (LogicalPhase &phase : partition.phases) { - for (LogicalStage &stage : phase.stages) { - stage.localSimtScopeCount = 0; - stage.scopeInputTensorBytes = 0; - stage.scopeOutputTensorBytes = 0; - auto merged = mergeSimtStageAnchors(anchorPlan, stage.simtAnchorIndices); - if (!merged) - continue; - { - const SimtAnchorDescriptor &anchor = *merged; - llvm::SmallVector roots; - const bool isRange = anchor.scopeOperations.size() > 1; - if (isRange) - llvm::append_range(roots, anchor.scopeOperations); - else - roots.push_back(anchor.operation); - - llvm::DenseSet inside; - for (Operation *root : roots) { - if (!root) - continue; - inside.insert(root); - root->walk([&](Operation *nested) { inside.insert(nested); }); - } + for (LogicalStage &stage : partition.stages) { + stage.localSimtScopeCount = 0; + stage.scopeInputTensorBytes = 0; + stage.scopeOutputTensorBytes = 0; + auto merged = mergeSimtStageAnchors(anchorPlan, stage.simtAnchorIndices); + if (!merged) + continue; + { + const SimtAnchorDescriptor &anchor = *merged; + llvm::SmallVector roots; + const bool isRange = anchor.scopeOperations.size() > 1; + if (isRange) + llvm::append_range(roots, anchor.scopeOperations); + else + roots.push_back(anchor.operation); + + llvm::DenseSet inside; + for (Operation *root : roots) { + if (!root) + continue; + inside.insert(root); + root->walk([&](Operation *nested) { inside.insert(nested); }); + } - llvm::SetVector captured; - for (Operation *operation : inside) - for (Value operand : operation->getOperands()) - if (!isValueDefinedInside(operand, inside)) - captured.insert(operand); - - llvm::SetVector returned; - for (Operation *root : roots) { - if (!root) - continue; - for (Value result : root->getResults()) { - // A single-op scope returns every result. A range scope mirrors - // wrapAnchorRange and returns only values with an outside user. - if (!isRange || - llvm::any_of(result.getUsers(), [&](Operation *user) { - return !inside.contains(user); - })) - returned.insert(result); - } - } + llvm::SetVector captured; + for (Operation *operation : inside) + for (Value operand : operation->getOperands()) + if (!isValueDefinedInside(operand, inside)) + captured.insert(operand); - // TritonToUnstructure cannot reconstruct offset information for a - // tensor-of-pointer returned by scope.scope. Capturing pointers is - // legal (the scope is not isolated from above), but returning pointer - // state would make this local Mixed implementation fail after route - // selection. Reject that implementation before it is scored; the - // same Stage remains legal in a whole-kernel pure-SIMT route. - if (llvm::any_of(returned, [](Value value) { - return isPointerLikeType(value.getType()); - })) { - stage.localSimtMaterializable = false; - stage.localSimtFactors.clear(); - stage.simtAnchorIndices.clear(); + llvm::SetVector returned; + for (Operation *root : roots) { + if (!root) continue; + for (Value result : root->getResults()) { + // A single-op scope returns every result. A range scope mirrors + // wrapAnchorRange and returns only values with an outside user. + if (!isRange || llvm::any_of(result.getUsers(), [&](Operation *user) { + return !inside.contains(user); + })) + returned.insert(result); } + } - ++stage.localSimtScopeCount; - stage.scopeInputTensorBytes += - staticTensorBytes(captured.getArrayRef()); - stage.scopeOutputTensorBytes += - staticTensorBytes(returned.getArrayRef()); + // TritonToUnstructure cannot reconstruct offset information for a + // tensor-of-pointer returned by scope.scope. Capturing pointers is + // legal (the scope is not isolated from above), but returning pointer + // state would make this local Mixed implementation fail after route + // selection. Reject that implementation before it is scored; the + // same Stage remains legal in a whole-kernel pure-SIMT route. + if (llvm::any_of(returned, [](Value value) { + return isPointerLikeType(value.getType()); + })) { + stage.localSimtMaterializable = false; + stage.localSimtFactors.clear(); + stage.simtAnchorIndices.clear(); + continue; } + + ++stage.localSimtScopeCount; + stage.scopeInputTensorBytes += staticTensorBytes(captured.getArrayRef()); + stage.scopeOutputTensorBytes += staticTensorBytes(returned.getArrayRef()); } } } @@ -1054,8 +753,9 @@ ProgramStructureAnalysis::analyze(ModuleOp module, // Stage boundaries must describe the program that the selected route will // materialize, rather than treating the pre-materialization textual order // as immutable. Normalize each compound anchor to its planned insertion - // point before PhaseBoundaryAnalysis performs a serial cut. The operation - // objects are not mutated here; only the analysis view is reordered. + // point before StageBoundaryAnalysis performs its semantic cut. The + // operation objects are not mutated here; only the analysis view is + // reordered. for (const SimtAnchorDescriptor &anchor : anchorPlan.anchors) { if (!anchor.materializable || anchor.scopeOperations.size() < 2 || !anchor.scopeInsertionPoint) @@ -1093,186 +793,311 @@ ProgramStructureAnalysis::analyze(ModuleOp module, structure.rootOperations = std::move(reordered); } + return structure; +} + +static int semanticKindPriority(StageCostModelKind kind) { + switch (kind) { + case StageCostModelKind::AutoBlockifyDispatch: + case StageCostModelKind::AutoBlockifyLoop: + return 100; + case StageCostModelKind::LoopCarriedRecurrence: + return 90; + case StageCostModelKind::RowwiseReduction: + return 80; + case StageCostModelKind::CubeRoofline: + case StageCostModelKind::TinyCubeRoofline: + return 70; + case StageCostModelKind::IndirectScalarMemory: + case StageCostModelKind::IndirectGatherMemory: + return 60; + case StageCostModelKind::IndependentPipelinedLoop: + return 50; + case StageCostModelKind::ConversionPack: + return 40; + case StageCostModelKind::ContinuousTileMemory: + case StageCostModelKind::ContinuousTileStore: + case StageCostModelKind::ContinuousShortLoad: + case StageCostModelKind::CachePolicyStore: + return 30; + case StageCostModelKind::PredicateMask: + case StageCostModelKind::LoopPredicate: + return 20; + case StageCostModelKind::IndexGeneration: + return 10; + default: + return 0; + } +} + +/// Scalar/index/predicate work is a supporting resource of a Stage, rather +/// than necessarily a Stage boundary of its own. StageCostModel accounts for +/// those resources inside every dominant model. Keeping this distinction +/// here prevents one Triton statement such as a masked load from being split +/// into artificial index, predicate and memory Stages. +static bool isSupportingSemanticKind(StageCostModelKind kind) { + switch (kind) { + case StageCostModelKind::ScalarIssue: + case StageCostModelKind::ScalarControl: + case StageCostModelKind::ScalarMath: + case StageCostModelKind::IndexGeneration: + case StageCostModelKind::PredicateMask: + case StageCostModelKind::LoopPredicate: + return true; + default: + return false; + } +} + +struct SourceStatement { + StringAttr file; + unsigned line = 0; + + explicit operator bool() const { return file && line != 0; } +}; + +static SourceStatement getSourceStatement(Location location) { + if (auto file = dyn_cast(location)) + return {file.getFilename(), file.getLine()}; + if (auto name = dyn_cast(location)) + return getSourceStatement(name.getChildLoc()); + if (auto callsite = dyn_cast(location)) { + SourceStatement callee = getSourceStatement(callsite.getCallee()); + return callee ? callee : getSourceStatement(callsite.getCaller()); + } + if (auto fused = dyn_cast(location)) + for (Location child : fused.getLocations()) { + SourceStatement statement = getSourceStatement(child); + if (statement) + return statement; + } + return {}; +} + +static bool haveSameSourceStatement(Operation *left, Operation *right) { + if (!left || !right) + return false; + SourceStatement lhs = getSourceStatement(left->getLoc()); + SourceStatement rhs = getSourceStatement(right->getLoc()); + return lhs && rhs && lhs.file == rhs.file && lhs.line == rhs.line; +} + +static llvm::Expected> +buildAnchorGroups(const ProgramStructure &structure, + const SimtAnchorPlan &anchorPlan) { + std::vector groups(structure.rootOperations.size(), -1); + int64_t nextGroup = 0; for (const SimtAnchorDescriptor &anchor : anchorPlan.anchors) { if (!anchor.materializable) continue; - auto addRoot = [&](Operation *operation) { + llvm::SmallVector positions; + auto addPosition = [&](Operation *operation) { Operation *root = getTopLevelSemanticRoot(operation); - if (root && llvm::is_contained(structure.rootOperations, root) && - !llvm::is_contained(structure.localSimtAnchorRoots, root)) - structure.localSimtAnchorRoots.push_back(root); + auto iterator = llvm::find(structure.rootOperations, root); + if (iterator == structure.rootOperations.end()) + return; + const size_t position = + static_cast(iterator - structure.rootOperations.begin()); + if (!llvm::is_contained(positions, position)) + positions.push_back(position); }; if (anchor.scopeOperations.empty()) - addRoot(anchor.operation); + addPosition(anchor.operation); else for (Operation *operation : anchor.scopeOperations) - addRoot(operation); - } - return structure; -} + addPosition(operation); + if (positions.empty()) + continue; + llvm::sort(positions); + for (size_t index = 1; index < positions.size(); ++index) + if (positions[index] != positions[index - 1] + 1) + return llvm::createStringError( + std::errc::invalid_argument, + "compound SIMT anchor roots are not contiguous after " + "ProgramStructureAnalysis"); -static std::optional -identifyPhaseBoundary(const SimdSimtFeatureSummary &features, - const StagePartitionerOptions &options) { - if (features.simtAnchors.triangularSolves.size() == 1 && - features.simtAnchors.count > 0) { - PhaseBoundaryPlan plan{PhaseBoundaryDomain::TriangularRecurrence, - "triangular_recurrence", - features.simtAnchors.triangularSolves.front()}; - return plan; - } - if (features.dotOps == 0 && features.reduceOps > 0 && - features.loadedIndexDependentMemoryOps > 0 && features.loadOps > 0 && - features.storeOps > 0) { - PhaseBoundaryPlan plan{PhaseBoundaryDomain::LoadedIndexRowwiseReduction, - "loaded_index_rowwise_reduction", std::nullopt}; - return plan; - } - if (features.dotOps > 0 && features.reduceOps == 0 && - features.loadedIndexDependentMemoryOps > 0 && - features.dotFlops <= options.tinyDotFlopsMax && features.loadOps > 0 && - features.storeOps > 0) { - PhaseBoundaryPlan plan{PhaseBoundaryDomain::IndirectUnderfilledDot, - "indirect_underfilled_dot", std::nullopt}; - return plan; + int64_t group = -1; + for (size_t position : positions) { + if (groups[position] < 0) + continue; + if (group < 0) + group = groups[position]; + else if (group != groups[position]) + return llvm::createStringError( + std::errc::invalid_argument, + "overlapping SIMT anchors define incompatible Stage boundaries"); + } + if (group < 0) + group = nextGroup++; + for (size_t position : positions) + groups[position] = group; } - return std::optional{}; -} - -llvm::Expected> -PhaseBoundaryAnalysis::analyze(ModuleOp module, - const SimtAnchorPlan &anchorPlan, - const SimdSimtFeatureSummary &features, - const StagePartitionerOptions &options) const { - auto plan = identifyPhaseBoundary(features, options); - if (!plan) - return std::optional{}; - auto structure = ProgramStructureAnalysis().analyze(module, anchorPlan); - if (!structure) - return structure.takeError(); - plan->rootOperations = std::move(structure->rootOperations); - plan->localSimtAnchorRoots = std::move(structure->localSimtAnchorRoots); - if (llvm::Error error = assignRootPhaseIds(*plan)) - return std::move(error); - return std::optional{std::move(*plan)}; + return groups; } llvm::Expected -StageBoundaryAnalysis::analyze(const PhaseBoundaryPlan &phasePlan, - const SimdSimtFeatureSummary &features, - const SimtAnchorPlan *anchorPlan) const { - if (!phasePlan.hasOperationGraph() || !anchorPlan) +StageBoundaryAnalysis::analyze(const ProgramStructure &structure, + const SimtAnchorPlan &anchorPlan) const { + if (structure.rootOperations.empty()) return llvm::createStringError( std::errc::invalid_argument, - "StageBoundaryAnalysis requires PreparedTTIR ownership"); + "StageBoundaryAnalysis requires ordered semantic roots"); + auto anchorGroups = buildAnchorGroups(structure, anchorPlan); + if (!anchorGroups) + return anchorGroups.takeError(); + StagePartition partition; - switch (phasePlan.domain) { - case PhaseBoundaryDomain::TriangularRecurrence: - if (!phasePlan.triangularSolve) + llvm::DenseSet owned; + for (size_t index = 0; index < structure.rootOperations.size();) { + Operation *root = structure.rootOperations[index]; + if (!root || !owned.insert(root).second) return llvm::createStringError( std::errc::invalid_argument, - "triangular PhaseBoundaryPlan has no recurrence facts"); - partition = - partitionTriangular(features, *phasePlan.triangularSolve, &phasePlan); - break; - case PhaseBoundaryDomain::LoadedIndexRowwiseReduction: - partition = partitionRowwise(features, &phasePlan); - break; - case PhaseBoundaryDomain::IndirectUnderfilledDot: - partition = partitionIndirectDot(features, &phasePlan); - break; + "StageBoundaryAnalysis received duplicate or null semantic root"); + + const int64_t anchorGroup = (*anchorGroups)[index]; + StageCostModelKind kind = classifySemanticRoot(root); + StageScheduleKind schedule = scheduleForSemanticRoot(root, kind); + LogicalStage stage; + stage.operations.push_back(root); + stage.iterationCount = semanticRootIterationCount(root); + stage.localSimtMaterializable = anchorGroup >= 0; + if (stage.localSimtMaterializable) + stage.localSimtFactors = {1}; + + size_t next = index + 1; + while (next < structure.rootOperations.size()) { + Operation *candidate = structure.rootOperations[next]; + const int64_t candidateAnchorGroup = (*anchorGroups)[next]; + const StageCostModelKind candidateKind = classifySemanticRoot(candidate); + const StageScheduleKind candidateSchedule = + scheduleForSemanticRoot(candidate, candidateKind); + const bool sameCompoundAnchor = + anchorGroup >= 0 && candidateAnchorGroup == anchorGroup; + const bool mergePlainStage = + anchorGroup < 0 && candidateAnchorGroup < 0 && + ((candidateKind == kind && candidateSchedule == schedule) || + (haveSameSourceStatement(stage.operations.back(), candidate) && + (isSupportingSemanticKind(kind) || + isSupportingSemanticKind(candidateKind)))) && + !hasOrderedStageBoundary(stage.operations.back()) && + !hasOrderedStageBoundary(candidate); + if (!sameCompoundAnchor && !mergePlainStage) + break; + if (!candidate || !owned.insert(candidate).second) + return llvm::createStringError( + std::errc::invalid_argument, + "StageBoundaryAnalysis overlaps semantic root ownership"); + stage.operations.push_back(candidate); + stage.iterationCount = + std::max(stage.iterationCount, semanticRootIterationCount(candidate)); + if (semanticKindPriority(candidateKind) > semanticKindPriority(kind)) { + kind = candidateKind; + schedule = candidateSchedule; + } + ++next; + } + stage.costModelKind = kind; + stage.scheduleKind = schedule; + stage.id = makeStageId(partition.stages.size(), kind); + partition.stages.push_back(std::move(stage)); + index = next; } - partition.domain = phasePlan.domainName; - if (llvm::Error error = - attachCompleteOperationOwnership(partition, phasePlan)) - return std::move(error); - attachExactAnchorOwnership(partition, *anchorPlan); + partition.operationOwnershipComplete = true; + partition.modeledOperationCount = + static_cast(structure.rootOperations.size()); + if (!partition.stages.empty()) + partition.stages.front().workload.paysKernelSetup = true; + + attachExactAnchorOwnership(partition, anchorPlan); deriveStageLiveValues(partition); - deriveLocalSimtScopeTraffic(partition, *anchorPlan); + deriveLocalSimtScopeTraffic(partition, anchorPlan); return partition; } llvm::Error StageFeatureAnalysis::analyze(StagePartition &partition) const { - for (LogicalPhase &phase : partition.phases) { - for (LogicalStage &stage : phase.stages) { - StageModelFeatures &facts = stage.features; - const double activeLaneRatio = facts.activeLaneRatio; - facts = StageModelFeatures{}; - facts.activeLaneRatio = activeLaneRatio; - llvm::DenseSet owned; + for (LogicalStage &stage : partition.stages) { + StageModelFeatures &facts = stage.features; + const double activeLaneRatio = facts.activeLaneRatio; + facts = StageModelFeatures{}; + facts.activeLaneRatio = activeLaneRatio; + llvm::DenseSet owned; + for (Operation *root : stage.operations) + collectOwnedOperationTree(root, owned); + bool hasMemory = false; + int64_t algorithmLoopCount = 0; + if (stage.costModelKind != StageCostModelKind::AutoBlockifyDispatch && + stage.costModelKind != StageCostModelKind::AutoBlockifyLoop) for (Operation *root : stage.operations) - collectOwnedOperationTree(root, owned); - bool hasMemory = false; - int64_t algorithmLoopCount = 0; - for (Operation *operation : owned) { - const llvm::StringRef name = operation->getName().getStringRef(); - if (name == "scf.for" || name == "scf.while") { - facts.hasLoop = true; - ++facts.loopBackedgeCount; - if (!operation->hasAttr("ta.auto_blockify_v1.loop")) - ++algorithmLoopCount; - if (!operation->hasAttr("ta.auto_blockify_v1.loop") && - operation->getNumRegions() > 0 && - !operation->getRegion(0).empty()) { - Block &body = operation->getRegion(0).front(); - const unsigned firstCarriedArgument = name == "scf.for" ? 1 : 0; - for (unsigned argumentIndex = firstCarriedArgument; - argumentIndex < body.getNumArguments(); ++argumentIndex) { - BlockArgument argument = body.getArgument(argumentIndex); - if (argument.use_empty()) - continue; - if (isPointerLikeType(argument.getType()) || - isAddressOnlyLoopValue(argument)) - facts.hasPointerInduction = true; - else - facts.hasLoopCarriedDataDependency = true; - } - if (name == "scf.for" && body.getNumArguments() > 0 && - isAddressOnlyLoopValue(body.getArgument(0))) + facts.replicatedByLocalSuperBlock |= isInsideAutoBlockifyV1Loop(root); + for (Operation *operation : owned) { + const llvm::StringRef name = operation->getName().getStringRef(); + if (name == "scf.for" || name == "scf.while") { + facts.hasLoop = true; + ++facts.loopBackedgeCount; + if (!operation->hasAttr("ta.auto_blockify_v1.loop")) + ++algorithmLoopCount; + if (!operation->hasAttr("ta.auto_blockify_v1.loop") && + operation->getNumRegions() > 0 && + !operation->getRegion(0).empty()) { + Block &body = operation->getRegion(0).front(); + const unsigned firstCarriedArgument = name == "scf.for" ? 1 : 0; + for (unsigned argumentIndex = firstCarriedArgument; + argumentIndex < body.getNumArguments(); ++argumentIndex) { + BlockArgument argument = body.getArgument(argumentIndex); + if (argument.use_empty()) + continue; + if (isPointerLikeType(argument.getType()) || + isAddressOnlyLoopValue(argument)) facts.hasPointerInduction = true; + else + facts.hasLoopCarriedDataDependency = true; } + if (name == "scf.for" && body.getNumArguments() > 0 && + isAddressOnlyLoopValue(body.getArgument(0))) + facts.hasPointerInduction = true; } - if (name == "scf.if" || name == "cf.cond_br") { - ++facts.conditionalBranchCount; - ++facts.divergentBranchCount; - } - if (name.contains("barrier") || name.contains("sync")) - ++facts.synchronizationCount; - if (name == "tt.load" || name == "tt.store" || name == "tt.gather" || - name.starts_with("tt.atomic")) { - hasMemory = true; - facts.hasIndirectMemory |= - isLoadedIndexDependentMemoryOp(operation) || - name == "tt.gather" || name.starts_with("tt.atomic"); - } - facts.hasReduction |= - name == "tt.reduce" || name == "tt.scan" || name == "linalg.reduce"; - facts.hasDot |= name == "tt.dot" || name.contains("matmul") || - name.contains("mmad"); - facts.hasConversionPack |= - name == "arith.extf" || name == "arith.truncf" || - name == "arith.fptosi" || name == "arith.fptoui" || - name == "arith.sitofp" || name == "arith.uitofp" || - name == "tt.fp_to_fp" || name.contains("convert") || - name.contains("pack") || name.contains("unpack"); } - facts.hasContiguousMemory = hasMemory && !facts.hasIndirectMemory; - if (algorithmLoopCount > 0 && stage.iterationCount > 1) { - if (facts.hasLoopCarriedDataDependency) - facts.parallelRecurrenceGroupCount = algorithmLoopCount; - facts.loopBackedgeCount = 1; - facts.conditionalBranchCount = std::max( - facts.conditionalBranchCount > 0 ? 1 : 0, - facts.conditionalBranchCount / algorithmLoopCount); - facts.divergentBranchCount = - std::max(facts.divergentBranchCount > 0 ? 1 : 0, - facts.divergentBranchCount / algorithmLoopCount); + if (name == "scf.if" || name == "cf.cond_br") { + ++facts.conditionalBranchCount; + ++facts.divergentBranchCount; + } + if (name.contains("barrier") || name.contains("sync")) + ++facts.synchronizationCount; + if (name == "tt.load" || name == "tt.store" || name == "tt.gather" || + name.starts_with("tt.atomic")) { + hasMemory = true; + facts.hasIndirectMemory |= isLoadedIndexDependentMemoryOp(operation) || + name == "tt.gather" || + name.starts_with("tt.atomic"); } - if (!facts.isValid()) - return llvm::createStringError(std::errc::invalid_argument, - "Stage '%s' has invalid features", - stage.id.c_str()); + facts.hasReduction |= + name == "tt.reduce" || name == "tt.scan" || name == "linalg.reduce"; + facts.hasDot |= + name == "tt.dot" || name.contains("matmul") || name.contains("mmad"); + facts.hasConversionPack |= + name == "arith.extf" || name == "arith.truncf" || + name == "arith.fptosi" || name == "arith.fptoui" || + name == "arith.sitofp" || name == "arith.uitofp" || + name == "tt.fp_to_fp" || name.contains("convert") || + name.contains("pack") || name.contains("unpack"); } + facts.hasContiguousMemory = hasMemory && !facts.hasIndirectMemory; + if (algorithmLoopCount > 0 && stage.iterationCount > 1) { + if (facts.hasLoopCarriedDataDependency) + facts.parallelRecurrenceGroupCount = algorithmLoopCount; + facts.loopBackedgeCount = 1; + facts.conditionalBranchCount = + std::max(facts.conditionalBranchCount > 0 ? 1 : 0, + facts.conditionalBranchCount / algorithmLoopCount); + facts.divergentBranchCount = + std::max(facts.divergentBranchCount > 0 ? 1 : 0, + facts.divergentBranchCount / algorithmLoopCount); + } + if (!facts.isValid()) + return llvm::createStringError(std::errc::invalid_argument, + "Stage '%s' has invalid features", + stage.id.c_str()); } return llvm::Error::success(); } @@ -1307,57 +1132,65 @@ llvm::Error StageKindClassifier::analyze(StagePartition &partition, return true; } }; - for (LogicalPhase &phase : partition.phases) { - for (LogicalStage &stage : phase.stages) { - const StageModelFeatures &facts = stage.features; - if (facts.hasDot && (facts.hasReduction || facts.hasIndirectMemory || - facts.hasLoopCarriedDataDependency)) - return llvm::createStringError( - std::errc::invalid_argument, - "requires_split: Stage '%s' owns incompatible dominant structures", - stage.id.c_str()); - - auto derive = [&]() { - if (facts.hasDot) - return stage.workload.dotFlops * stage.iterationCount <= - static_cast( - std::max(1, tinyDotFlopsMax)) - ? StageCostModelKind::TinyCubeRoofline - : StageCostModelKind::CubeRoofline; - if (facts.hasReduction) - return StageCostModelKind::RowwiseReduction; - if (facts.hasConversionPack) - return StageCostModelKind::ConversionPack; - if (facts.hasLoop) - return facts.hasLoopCarriedDataDependency - ? StageCostModelKind::LoopCarriedRecurrence - : StageCostModelKind::IndependentPipelinedLoop; - if (facts.hasIndirectMemory) - return StageCostModelKind::IndirectGatherMemory; - if (facts.hasContiguousMemory) - return stage.workload.storeBytes > 0.0 && - stage.workload.loadBytes == 0.0 - ? StageCostModelKind::ContinuousTileStore - : StageCostModelKind::ContinuousTileMemory; - return StageCostModelKind::ScalarIssue; - }; - if (!compatible(stage.costModelKind, facts)) - stage.costModelKind = derive(); - if (!compatible(stage.costModelKind, facts) || - (stage.costModelKind == StageCostModelKind::TinyCubeRoofline && - stage.workload.dotFlops * stage.iterationCount > - static_cast(std::max(1, tinyDotFlopsMax)))) - return llvm::createStringError( - std::errc::invalid_argument, - "Stage '%s' operation graph does not match StageCostModelKind '%s'", - stage.id.c_str(), - stringifyStageCostModel(stage.costModelKind).str().c_str()); - if (stage.costModelKind == StageCostModelKind::IndependentPipelinedLoop) - stage.scheduleKind = StageScheduleKind::IndependentPipelined; - else if (stage.costModelKind == StageCostModelKind::LoopCarriedRecurrence) - stage.scheduleKind = StageScheduleKind::LoopCarriedSerial; - } + for (LogicalStage &stage : partition.stages) { + const StageModelFeatures &facts = stage.features; + if (stage.costModelKind == StageCostModelKind::AutoBlockifyDispatch || + stage.costModelKind == StageCostModelKind::AutoBlockifyLoop) + continue; + if (facts.hasDot && (facts.hasReduction || facts.hasIndirectMemory || + facts.hasLoopCarriedDataDependency)) + return llvm::createStringError( + std::errc::invalid_argument, + "requires_split: Stage '%s' owns incompatible dominant structures", + stage.id.c_str()); + + auto derive = [&]() { + if (facts.hasLoopCarriedDataDependency) + return StageCostModelKind::LoopCarriedRecurrence; + if (facts.hasReduction) + return StageCostModelKind::RowwiseReduction; + if (facts.hasDot) + return stage.workload.dotFlops * stage.iterationCount <= + static_cast( + std::max(1, tinyDotFlopsMax)) + ? StageCostModelKind::TinyCubeRoofline + : StageCostModelKind::CubeRoofline; + if (facts.hasLoop) + return StageCostModelKind::IndependentPipelinedLoop; + if (facts.hasIndirectMemory) + return StageCostModelKind::IndirectGatherMemory; + if (facts.hasConversionPack) + return StageCostModelKind::ConversionPack; + if (facts.hasContiguousMemory) + return stage.workload.storeBytes > 0.0 && + stage.workload.loadBytes == 0.0 + ? StageCostModelKind::ContinuousTileStore + : StageCostModelKind::ContinuousTileMemory; + return StageCostModelKind::ScalarIssue; + }; + const StageCostModelKind derived = derive(); + // Strong operation-graph semantics are authoritative. Scalar + // sub-kinds remain useful only when no dominant structure is present. + if (semanticKindPriority(derived) > 0 || + !compatible(stage.costModelKind, facts)) + stage.costModelKind = derived; + if (!compatible(stage.costModelKind, facts) || + (stage.costModelKind == StageCostModelKind::TinyCubeRoofline && + stage.workload.dotFlops * stage.iterationCount > + static_cast(std::max(1, tinyDotFlopsMax)))) + return llvm::createStringError( + std::errc::invalid_argument, + "Stage '%s' operation graph does not match StageCostModelKind '%s'", + stage.id.c_str(), + stringifyStageCostModel(stage.costModelKind).str().c_str()); + if (stage.costModelKind == StageCostModelKind::IndependentPipelinedLoop) + stage.scheduleKind = StageScheduleKind::IndependentPipelined; + else if (stage.costModelKind == StageCostModelKind::LoopCarriedRecurrence) + stage.scheduleKind = StageScheduleKind::LoopCarriedSerial; } + for (auto indexedStage : llvm::enumerate(partition.stages)) + indexedStage.value().id = + makeStageId(indexedStage.index(), indexedStage.value().costModelKind); return llvm::Error::success(); } @@ -1366,80 +1199,67 @@ llvm::Error StageWorkloadAnalysis::analyze(StagePartition &partition) const { return llvm::createStringError( std::errc::invalid_argument, "StageWorkloadAnalysis requires complete operation ownership"); - for (LogicalPhase &phase : partition.phases) { - for (LogicalStage &stage : phase.stages) { - StageWorkload work; - work.paysKernelSetup = stage.workload.paysKernelSetup; - const int64_t loopCount = countAlgorithmLoops(stage); - const int64_t fallbackLoopTripCount = - loopCount > 0 ? std::max(1, stage.iterationCount / loopCount) - : 1; - for (Operation *root : stage.operations) - accumulateDynamicOperationTree(root, work, 1.0, fallbackLoopTripCount); - recomputeIssueElements(work); - stage.workload = std::move(work); - makePerIteration(stage); - if (!stage.workload.isFiniteAndNonNegative()) - return llvm::createStringError( - std::errc::invalid_argument, - "Stage '%s' has invalid operation-derived workload", - stage.id.c_str()); - } + for (LogicalStage &stage : partition.stages) { + StageWorkload work; + work.paysKernelSetup = stage.workload.paysKernelSetup; + const int64_t loopCount = countAlgorithmLoops(stage); + const int64_t fallbackLoopTripCount = + loopCount > 0 ? std::max(1, stage.iterationCount / loopCount) + : 1; + for (Operation *root : stage.operations) + accumulateDynamicOperationTree(root, work, 1.0, fallbackLoopTripCount); + recomputeIssueElements(work); + stage.workload = std::move(work); + makePerIteration(stage); + if (!stage.workload.isFiniteAndNonNegative()) + return llvm::createStringError( + std::errc::invalid_argument, + "Stage '%s' has invalid operation-derived workload", + stage.id.c_str()); } return llvm::Error::success(); } llvm::Error StagePartitionVerifier::verify(const StagePartition &partition) const { - if (partition.phases.empty()) + if (partition.stages.empty()) return llvm::createStringError(std::errc::invalid_argument, - "StagePartition has no Phase"); - llvm::StringSet<> phaseIds; + "StagePartition has no Stage"); llvm::StringSet<> stageIds; llvm::DenseSet ownedOperations; llvm::DenseSet ownedAnchors; - for (const LogicalPhase &phase : partition.phases) { - if (phase.id.empty() || !phaseIds.insert(phase.id).second) + for (const LogicalStage &stage : partition.stages) { + if (stage.id.empty() || !stageIds.insert(stage.id).second) return llvm::createStringError(std::errc::invalid_argument, - "StagePartition has duplicate Phase id"); - if (phase.stages.empty()) + "StagePartition has duplicate Stage id"); + if (stage.iterationCount < 1) return llvm::createStringError(std::errc::invalid_argument, - "Phase '%s' has no Stage", - phase.id.c_str()); - for (const LogicalStage &stage : phase.stages) { - if (stage.id.empty() || !stageIds.insert(stage.id).second) - return llvm::createStringError(std::errc::invalid_argument, - "StagePartition has duplicate Stage id"); - if (stage.iterationCount < 1) - return llvm::createStringError(std::errc::invalid_argument, - "Stage '%s' has invalid iteration count", - stage.id.c_str()); - if (stage.localSimtMaterializable && - partition.operationOwnershipComplete && stage.operations.empty()) - return llvm::createStringError( - std::errc::invalid_argument, - "materializable Stage '%s' has no operation ownership", - stage.id.c_str()); - if (stage.localSimtMaterializable && - partition.operationOwnershipComplete && - stage.simtAnchorIndices.empty()) - return llvm::createStringError( - std::errc::invalid_argument, - "materializable Stage '%s' has no exact SIMT anchor ownership", - stage.id.c_str()); - if (partition.operationOwnershipComplete) - for (unsigned index : stage.simtAnchorIndices) - if (!ownedAnchors.insert(index).second) - return llvm::createStringError( - std::errc::invalid_argument, - "StagePartition SIMT anchor ownership overlaps"); - if (partition.operationOwnershipComplete) - for (Operation *operation : stage.operations) - if (!operation || !ownedOperations.insert(operation).second) - return llvm::createStringError( - std::errc::invalid_argument, - "StagePartition operation ownership overlaps"); - } + "Stage '%s' has invalid iteration count", + stage.id.c_str()); + if (stage.localSimtMaterializable && partition.operationOwnershipComplete && + stage.operations.empty()) + return llvm::createStringError( + std::errc::invalid_argument, + "materializable Stage '%s' has no operation ownership", + stage.id.c_str()); + if (stage.localSimtMaterializable && partition.operationOwnershipComplete && + stage.simtAnchorIndices.empty()) + return llvm::createStringError( + std::errc::invalid_argument, + "materializable Stage '%s' has no exact SIMT anchor ownership", + stage.id.c_str()); + if (partition.operationOwnershipComplete) + for (unsigned index : stage.simtAnchorIndices) + if (!ownedAnchors.insert(index).second) + return llvm::createStringError( + std::errc::invalid_argument, + "StagePartition SIMT anchor ownership overlaps"); + if (partition.operationOwnershipComplete) + for (Operation *operation : stage.operations) + if (!operation || !ownedOperations.insert(operation).second) + return llvm::createStringError( + std::errc::invalid_argument, + "StagePartition operation ownership overlaps"); } if (partition.operationOwnershipComplete && static_cast(ownedOperations.size()) != @@ -1464,57 +1284,50 @@ StageModeLegalityAnalysis::analyze(StagePartition &partition, // already reduced the target maximum to F2 for num_warps=32 (or to F1 for // a smaller runtime grid). const int64_t localMaximum = scopeSuperblockMaterializable ? maximum : 1; - for (LogicalPhase &phase : partition.phases) { - for (LogicalStage &stage : phase.stages) { - stage.simdLegal = true; - stage.simtLegal = true; - stage.legalSimtFactors = {1}; - // A pure-SIMT SuperBlock factor is a whole-kernel schedule, not a - // recurrence-only annotation. Every SIMT Stage must therefore expose - // the same factor candidates; KernelRouteSolver keeps the chosen factor - // uniform. Local mixed scopes stay restricted by localSimtFactors - // (F1 unless Scope SuperBlock materialization is explicitly available). - if (maximum >= 2) - stage.legalSimtFactors.push_back(2); - if (maximum >= 4) - stage.legalSimtFactors.push_back(4); - if (stage.localSimtMaterializable) { - // The ABI-v2 scope materializer batches complete logical programs - // around this Stage. F2/F4 therefore does not require multiple - // recurrence groups inside one logical program; that older W2/W4 - // interpretation was only warp widening, not a SuperBlock. - stage.localSimtFactors = {1}; - if (scopeSuperblockMaterializable) - for (int64_t factor : {2, 4}) - if (factor <= localMaximum) - stage.localSimtFactors.push_back(factor); - } - if (stage.localSimtMaterializable && - (stage.localSimtFactors.empty() || - llvm::any_of(stage.localSimtFactors, [&](int64_t factor) { - return factor < 1 || factor > localMaximum || - (factor != 1 && factor != 2 && factor != 4); - }))) - return llvm::createStringError( - std::errc::invalid_argument, - "local SIMT factors are invalid for Stage '%s'", stage.id.c_str()); + for (LogicalStage &stage : partition.stages) { + stage.simdLegal = true; + stage.simtLegal = true; + stage.legalSimtFactors = {1}; + // A pure-SIMT SuperBlock factor is a whole-kernel schedule, not a + // recurrence-only annotation. Every SIMT Stage must therefore expose + // the same factor candidates; KernelRouteSolver keeps the chosen factor + // uniform. Local mixed scopes stay restricted by localSimtFactors + // (F1 unless Scope SuperBlock materialization is explicitly available). + if (maximum >= 2) + stage.legalSimtFactors.push_back(2); + if (maximum >= 4) + stage.legalSimtFactors.push_back(4); + if (stage.localSimtMaterializable) { + // The ABI-v2 scope materializer batches complete logical programs + // around this Stage. F2/F4 therefore does not require multiple + // recurrence groups inside one logical program; that older W2/W4 + // interpretation was only warp widening, not a SuperBlock. + stage.localSimtFactors = {1}; + if (scopeSuperblockMaterializable && stage.localSuperblockMaterializable) + for (int64_t factor : {2, 4}) + if (factor <= localMaximum) + stage.localSimtFactors.push_back(factor); } + if (stage.localSimtMaterializable && + (stage.localSimtFactors.empty() || + llvm::any_of(stage.localSimtFactors, [&](int64_t factor) { + return factor < 1 || factor > localMaximum || + (factor != 1 && factor != 2 && factor != 4); + }))) + return llvm::createStringError( + std::errc::invalid_argument, + "local SIMT factors are invalid for Stage '%s'", stage.id.c_str()); } return llvm::Error::success(); } -llvm::Expected> +llvm::Expected StagePartitioner::partition(ModuleOp module, const SimtAnchorPlan &anchorPlan, - const SimdSimtFeatureSummary &features, const StagePartitionerOptions &options) const { - auto phasePlan = - PhaseBoundaryAnalysis().analyze(module, anchorPlan, features, options); - if (!phasePlan) - return phasePlan.takeError(); - if (!*phasePlan) - return std::optional{}; - auto result = - StageBoundaryAnalysis().analyze(**phasePlan, features, &anchorPlan); + auto structure = ProgramStructureAnalysis().analyze(module, anchorPlan); + if (!structure) + return structure.takeError(); + auto result = StageBoundaryAnalysis().analyze(*structure, anchorPlan); if (!result) return result.takeError(); StageWorkloadAnalysis workloadAnalysis; @@ -1533,5 +1346,5 @@ StagePartitioner::partition(ModuleOp module, const SimtAnchorPlan &anchorPlan, return std::move(error); if (llvm::Error error = StagePartitionVerifier().verify(*result)) return std::move(error); - return std::optional{std::move(*result)}; + return std::move(*result); } diff --git a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/SimdSimtCostModel.cpp b/third_party/ascend/costmodel/lib/AscendModel/RouteModel/SimdSimtCostModel.cpp index 9e681d8240..ebe60f6759 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/SimdSimtCostModel.cpp +++ b/third_party/ascend/costmodel/lib/AscendModel/RouteModel/SimdSimtCostModel.cpp @@ -16,8 +16,6 @@ #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Operation.h" #include "llvm/ADT/ArrayRef.h" -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/DenseSet.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallString.h" #include "llvm/ADT/SmallVector.h" @@ -32,7 +30,6 @@ #include #include #include -#include #include #include #include @@ -127,70 +124,6 @@ class ProfileJSONReader { std::string error; }; -static std::optional getConstantInteger(Value value) { - Operation *definingOp = value.getDefiningOp(); - if (!definingOp || definingOp->getName().getStringRef() != "arith.constant") - return std::nullopt; - if (auto integer = definingOp->getAttrOfType("value")) - return integer.getInt(); - return std::nullopt; -} - -static std::optional getKnownStaticLoopTripCount(Operation *op) { - if (!op || op->getName().getStringRef() != "scf.for" || - op->getNumOperands() < 3) - return std::nullopt; - auto lower = getConstantInteger(op->getOperand(0)); - auto upper = getConstantInteger(op->getOperand(1)); - auto step = getConstantInteger(op->getOperand(2)); - if (!lower || !upper || !step || *step == 0) - return std::nullopt; - int64_t span = *upper - *lower; - if (span > 0 && *step > 0) - return std::max(1, (span + *step - 1) / *step); - if (span < 0 && *step < 0) { - int64_t positiveSpan = -span; - int64_t positiveStep = -*step; - return std::max(1, - (positiveSpan + positiveStep - 1) / positiveStep); - } - return std::nullopt; -} - -static int64_t getModeledLoopTripCount( - Operation *op, - const llvm::DenseMap &structuralTripEstimates) { - if (auto knownTripCount = getKnownStaticLoopTripCount(op)) - return *knownTripCount; - if (auto iterator = structuralTripEstimates.find(op); - iterator != structuralTripEstimates.end()) - return iterator->second; - return 1; -} - -static int64_t getLoopMultiplier( - Operation *op, - const llvm::DenseMap &structuralTripEstimates) { - int64_t multiplier = 1; - for (Operation *parent = op->getParentOp(); parent; - parent = parent->getParentOp()) { - if (parent->getName().getStringRef() != "scf.for") - continue; - // AutoBlockify V1 wraps the original logical program in a physical-core - // scheduling loop. Its dispatch cost is modeled as a separate phase; - // it must not multiply every algorithm operation as if it were an - // algorithmic loop. - if (parent->hasAttr("ta.auto_blockify_v1.schedule")) - continue; - int64_t tripCount = - getModeledLoopTripCount(parent, structuralTripEstimates); - if (tripCount > 0 && - multiplier <= std::numeric_limits::max() / tripCount) - multiplier *= tripCount; - } - return multiplier; -} - static double resolveNumberOrMeasurement( const llvm::json::Object &object, llvm::StringRef numberKey, llvm::StringRef measurementKey, llvm::StringRef expectedUnit, @@ -576,7 +509,7 @@ loadCandidateProfile(llvm::StringRef requestedPath) { return profile; } -static llvm::Expected> evaluateStageModel( +static llvm::Expected evaluateStageModel( const SimdSimtFeatureSummary &features, const CandidateProfile &profile, unsigned numWarps, bool wholeKernelSuperblockMaterializable, bool scopeSuperblockMaterializable, int64_t logicalProgramCountHint, @@ -609,16 +542,14 @@ static llvm::Expected> evaluateStageModel( "Stage model requires PreparedTTIR and " "its anchor plan"); auto partition = - partitioner.partition(module, *anchorPlan, features, partitionerOptions); + partitioner.partition(module, *anchorPlan, partitionerOptions); if (!partition) return partition.takeError(); - if (!*partition) - return std::optional{}; HardwareProfile hardwareProfile = profile.hardware; hardwareProfile.logicalWarpGroupCount = std::max(1, numWarps); StageCostEvaluator evaluator; - auto costTable = evaluator.evaluate(**partition, hardwareProfile); + auto costTable = evaluator.evaluate(*partition, hardwareProfile); if (!costTable) return costTable.takeError(); costTable->logicalProgramCountHint = logicalProgramCountHint; @@ -626,7 +557,7 @@ static llvm::Expected> evaluateStageModel( auto routes = solveStageRoutes(*costTable, hardwareProfile.transition); if (!routes) return routes.takeError(); - return std::optional{std::move(*routes)}; + return std::move(*routes); } static llvm::SmallVector> @@ -698,9 +629,6 @@ toTriangularSolveFactsJSON(const TriangularSolveFacts &facts) { llvm::json::Object SimtAnchorFeatureSummary::toJSON() const { llvm::json::Object result; result["count"] = count; - result["conditional_branch_count"] = conditionalBranchCount; - result["divergent_branch_count"] = divergentBranchCount; - result["active_lane_ratio"] = activeLaneRatio; llvm::json::Array triangularFacts; for (const TriangularSolveFacts &facts : triangularSolves) triangularFacts.push_back(toTriangularSolveFactsJSON(facts)); @@ -711,19 +639,8 @@ llvm::json::Object SimtAnchorFeatureSummary::toJSON() const { llvm::json::Object SimdSimtFeatureSummary::toJSON() const { llvm::json::Object result; - result["load_ops"] = loadOps; - result["store_ops"] = storeOps; - result["reduce_ops"] = reduceOps; - result["dot_ops"] = dotOps; - result["loaded_index_dependent_memory_ops"] = loadedIndexDependentMemoryOps; - result["dot_flops"] = dotFlops; - result["static_loop_trip_count_max"] = staticLoopTripCountMax; - result["conditional_branch_count"] = conditionalBranchCount; - result["divergent_branch_count"] = divergentBranchCount; - result["active_lane_ratio"] = activeLaneRatio; llvm::json::Object postTransform; postTransform["auto_blockify_v1_applied"] = autoBlockifyV1Applied; - postTransform["auto_blockify_v1_loop_count"] = autoBlockifyV1LoopCount; result["post_transform"] = std::move(postTransform); result["has_explicit_scope"] = hasExplicitScope; result["simt_anchors"] = simtAnchors.toJSON(); @@ -802,8 +719,6 @@ mlir::ascend::analyzeSimdSimtFeatures(ModuleOp module, "cannot analyze a null ModuleOp"); SimdSimtFeatureSummary features; - llvm::DenseSet anchorRoots; - llvm::DenseMap structuralTripEstimates; features.simtAnchors.count = llvm::count_if( anchorPlan.anchors, [](const SimtAnchorDescriptor &anchor) { return anchor.materializable; }); @@ -812,75 +727,14 @@ mlir::ascend::analyzeSimdSimtFeatures(ModuleOp module, for (const SimtAnchorDescriptor &anchor : anchorPlan.anchors) { if (anchor.triangularSolve) features.simtAnchors.triangularSolves.push_back(*anchor.triangularSolve); - if (!anchor.materializable) - continue; - for (Operation *operation : anchor.scopeOperations) - if (operation) - anchorRoots.insert(operation); - if (anchor.scopeOperations.empty() && anchor.operation) - anchorRoots.insert(anchor.operation); - if (anchor.kind == SimtAnchorKind::TriangularSolveLoop) - for (Operation *operation : anchor.scopeOperations) - if (operation && operation->getName().getStringRef() == "scf.for") - structuralTripEstimates[operation] = 14; } - auto isInAnchor = [&](Operation *operation) { - for (; operation; operation = operation->getParentOp()) - if (anchorRoots.contains(operation)) - return true; - return false; - }; module.walk([&](Operation *operation) { - if (operation->hasAttr("ta.auto_blockify_v1")) + if (operation->hasAttr("ta.auto_blockify_v1") || + operation->hasAttr("ta.auto_blockify_v1.loop")) features.autoBlockifyV1Applied = true; - if (operation->hasAttr("ta.auto_blockify_v1.loop")) { - features.autoBlockifyV1Applied = true; - ++features.autoBlockifyV1LoopCount; - } - if (operation->hasAttr("ta.auto_blockify_v1.schedule")) - return; - - llvm::StringRef name = operation->getName().getStringRef(); - const bool inAnchor = isInAnchor(operation); - const int64_t multiplier = - getLoopMultiplier(operation, structuralTripEstimates); - features.hasExplicitScope |= name == "scope.scope"; - features.loadOps += name == "tt.load"; - features.storeOps += name == "tt.store"; - features.reduceOps += name == "tt.reduce"; - features.dotOps += name == "tt.dot"; - - if (name == "scf.if" || name == "cf.cond_br") { - ++features.conditionalBranchCount; - if (inAnchor) - ++features.simtAnchors.conditionalBranchCount; - const bool divergent = - operation->getNumOperands() > 0 && - isa(operation->getOperand(0).getType()); - features.divergentBranchCount += divergent; - if (inAnchor) - features.simtAnchors.divergentBranchCount += divergent; - } - if (name == "scf.for") { - const int64_t trip = - getModeledLoopTripCount(operation, structuralTripEstimates); - features.staticLoopTripCountMax = - std::max(features.staticLoopTripCountMax, trip); - } - if (isLoadedIndexDependentMemoryOp(operation)) - ++features.loadedIndexDependentMemoryOps; - if (name == "tt.dot" && operation->getNumOperands() >= 2) { - auto lhs = dyn_cast(operation->getOperand(0).getType()); - auto rhs = dyn_cast(operation->getOperand(1).getType()); - if (lhs && rhs && lhs.getRank() >= 2 && rhs.getRank() >= 2) { - const int64_t m = lhs.getShape()[lhs.getRank() - 2]; - const int64_t k = lhs.getShape()[lhs.getRank() - 1]; - const int64_t n = rhs.getShape()[rhs.getRank() - 1]; - if (m > 0 && n > 0 && k > 0) - features.dotFlops += 2 * m * n * k * multiplier; - } - } + features.hasExplicitScope |= + operation->getName().getStringRef() == "scope.scope"; }); return features; } @@ -927,33 +781,24 @@ estimateSimdSimtCandidatesImpl(const SimdSimtFeatureSummary &features, options.physicalVectorCoreCountHint, module, anchorPlan); if (!stageModel) return stageModel.takeError(); - if (*stageModel) { - report.stageModel = std::move(**stageModel); - report.candidateCosts.allSimd = report.stageModel.allSimd.totalCycles; - report.candidateCosts.allSimtOnly = report.stageModel.allSimt.totalCycles; - report.candidateCosts.mixedSimdSimt = report.stageModel.mixed.totalCycles; - report.allSimdCandidateLegal &= report.stageModel.allSimd.legal; - report.allSimtOnlyCandidateLegal &= report.stageModel.allSimt.legal; - report.mixedCandidateLegal &= report.stageModel.mixed.legal; - const unsigned legalCandidateCount = - static_cast(report.allSimdCandidateLegal) + - static_cast(report.allSimtOnlyCandidateLegal) + - static_cast(report.mixedCandidateLegal); - if (legalCandidateCount == 0) - return llvm::createStringError( - std::errc::not_supported, - "Stage Route Model found no materializable candidate"); - report.decision = chooseBest( - report.candidateCosts, report.allSimdCandidateLegal, - report.allSimtOnlyCandidateLegal, report.mixedCandidateLegal); - return report; - } - - // Unknown Stage domains are deliberately not scored. The online Route - // Model has no legacy whole-kernel fallback; the selector leaves the - // existing backend-default lowering unchanged. - report.stageModel.applied = false; - report.unsupported.push_back("stage_model_not_applicable"); + report.stageModel = std::move(*stageModel); + report.candidateCosts.allSimd = report.stageModel.allSimd.totalCycles; + report.candidateCosts.allSimtOnly = report.stageModel.allSimt.totalCycles; + report.candidateCosts.mixedSimdSimt = report.stageModel.mixed.totalCycles; + report.allSimdCandidateLegal &= report.stageModel.allSimd.legal; + report.allSimtOnlyCandidateLegal &= report.stageModel.allSimt.legal; + report.mixedCandidateLegal &= report.stageModel.mixed.legal; + const unsigned legalCandidateCount = + static_cast(report.allSimdCandidateLegal) + + static_cast(report.allSimtOnlyCandidateLegal) + + static_cast(report.mixedCandidateLegal); + if (legalCandidateCount == 0) + return llvm::createStringError( + std::errc::not_supported, + "Stage Route Model found no materializable candidate"); + report.decision = + chooseBest(report.candidateCosts, report.allSimdCandidateLegal, + report.allSimtOnlyCandidateLegal, report.mixedCandidateLegal); return report; } diff --git a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageCostModels.cpp b/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageCostModels.cpp index ed5f7a7966..d693c892a1 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageCostModels.cpp +++ b/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageCostModels.cpp @@ -5,6 +5,7 @@ #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/StringSet.h" #include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/raw_ostream.h" #include #include @@ -21,6 +22,22 @@ static double iterations(const LogicalStage &stage) { return static_cast(std::max(1, stage.iterationCount)); } +static std::vector +collectSourceLocations(const LogicalStage &stage) { + std::vector result; + llvm::StringSet<> seen; + for (Operation *operation : stage.operations) { + std::string location; + llvm::raw_string_ostream stream(location); + operation->getLoc().print(stream); + stream.flush(); + if (location.empty() || !seen.insert(location).second) + continue; + result.push_back(std::move(location)); + } + return result; +} + static double controlBody(const StageResourceCycles &resources) { return resources.loopControl + resources.branchControl + resources.divergence + resources.synchronization; @@ -221,8 +238,18 @@ static double estimateStage(const LogicalStage &stage, controlBody(r) + r.spill, r.issue) : serialBody(r); - if (mode == StageMode::SIMD) - return r.setup + count * critical; + if (mode == StageMode::SIMD) { + // A loop-carried tensor is not ordinary embarrassingly-parallel vector + // work: the updated state must remain live until the next recurrence + // step. The operation-throughput terms above account for arithmetic, + // but not this persistent register/stack traffic. Charge the exact + // SSA live-out footprint once per Stage invocation using the target + // profile's persistent-state byte rate. + const double persistentState = + static_cast(stage.liveOutBytes) / + profile.superblockPersistentStateBytesPerCycle; + return r.setup + count * critical + persistentState; + } const int64_t groups = std::max( 1, std::min(stage.features.parallelRecurrenceGroupCount, profile.logicalWarpGroupCount)); @@ -376,104 +403,95 @@ bool HardwareProfile::isValid() const { llvm::Expected StageCostEvaluator::evaluate(const StagePartition &partition, const HardwareProfile &profile) const { - if (partition.domain.empty() || partition.phases.empty()) + if (partition.stages.empty()) return llvm::createStringError( std::errc::invalid_argument, - "StagePartition requires a domain and at least one Phase"); + "StagePartition requires at least one Stage"); if (!profile.isValid()) return llvm::createStringError(std::errc::invalid_argument, "HardwareProfile is invalid"); StageCostTable table; - table.domain = partition.domain; table.operationOwnershipComplete = partition.operationOwnershipComplete; table.modeledOperationCount = partition.modeledOperationCount; table.profileVersion = profile.profileVersion; llvm::StringSet<> stageIds; - for (const LogicalPhase &phase : partition.phases) { - if (phase.id.empty() || phase.stages.empty()) + for (const LogicalStage &stage : partition.stages) { + if (stage.id.empty() || !stageIds.insert(stage.id).second) + return llvm::createStringError( + std::errc::invalid_argument, + "Stage ids must be non-empty and unique: '%s'", stage.id.c_str()); + if (stage.iterationCount <= 0 || !stage.features.isValid() || + !stage.workload.isFiniteAndNonNegative()) + return llvm::createStringError( + std::errc::invalid_argument, + "Stage '%s' has invalid iteration/features", stage.id.c_str()); + if (!stage.simdLegal && !stage.simtLegal) return llvm::createStringError(std::errc::invalid_argument, - "every Phase requires an id and Stage"); - LogicalPhaseCost phaseCost; - phaseCost.id = phase.id; - - for (const LogicalStage &stage : phase.stages) { - if (stage.id.empty() || !stageIds.insert(stage.id).second) - return llvm::createStringError( - std::errc::invalid_argument, - "Stage ids must be non-empty and unique: '%s'", stage.id.c_str()); - if (stage.iterationCount <= 0 || !stage.features.isValid() || - !stage.workload.isFiniteAndNonNegative()) - return llvm::createStringError( - std::errc::invalid_argument, - "Stage '%s' has invalid iteration/features", stage.id.c_str()); - if (!stage.simdLegal && !stage.simtLegal) - return llvm::createStringError(std::errc::invalid_argument, - "Stage '%s' has no legal StageMode", - stage.id.c_str()); - if (stage.simtLegal && stage.legalSimtFactors.empty()) - return llvm::createStringError( - std::errc::invalid_argument, - "SIMT Stage '%s' has no legal SuperBlock factor", stage.id.c_str()); - - LogicalStageCost logicalCost; - logicalCost.id = stage.id; - logicalCost.model = stringifyStageCostModel(stage.costModelKind).str(); - logicalCost.schedule = stage.scheduleKind; - logicalCost.iterationCount = stage.iterationCount; - logicalCost.features = stage.features; - logicalCost.workload = stage.workload; - logicalCost.ownedOperationCount = - static_cast(stage.operations.size()); - logicalCost.liveInCount = static_cast(stage.liveIns.size()); - logicalCost.liveOutCount = static_cast(stage.liveOuts.size()); - logicalCost.liveInBytes = stage.liveInBytes; - logicalCost.liveOutBytes = stage.liveOutBytes; - logicalCost.localSimtScopeCount = stage.localSimtScopeCount; - logicalCost.scopeInputTensorBytes = stage.scopeInputTensorBytes; - logicalCost.scopeOutputTensorBytes = stage.scopeOutputTensorBytes; - logicalCost.simtAnchorIndices = stage.simtAnchorIndices; - logicalCost.localSimtMaterializable = stage.localSimtMaterializable; - logicalCost.legalSimtFactors = stage.legalSimtFactors; - logicalCost.localSimtFactors = stage.localSimtFactors; + "Stage '%s' has no legal StageMode", + stage.id.c_str()); + if (stage.simtLegal && stage.legalSimtFactors.empty()) + return llvm::createStringError( + std::errc::invalid_argument, + "SIMT Stage '%s' has no legal SuperBlock factor", stage.id.c_str()); - llvm::SmallVector implementations; - if (stage.simdLegal) - implementations.push_back({StageMode::SIMD, 1, false}); - if (stage.simtLegal) - for (int64_t factor : stage.legalSimtFactors) - implementations.push_back({StageMode::SIMT, factor, false}); - if (stage.simtLegal && stage.localSimtMaterializable) - for (int64_t factor : stage.localSimtFactors) - implementations.push_back({StageMode::SIMT, factor, true}); + LogicalStageCost logicalCost; + logicalCost.id = stage.id; + logicalCost.model = stringifyStageCostModel(stage.costModelKind).str(); + logicalCost.schedule = stage.scheduleKind; + logicalCost.iterationCount = stage.iterationCount; + logicalCost.features = stage.features; + logicalCost.workload = stage.workload; + logicalCost.ownedOperationCount = + static_cast(stage.operations.size()); + logicalCost.sourceLocations = collectSourceLocations(stage); + logicalCost.liveInCount = static_cast(stage.liveIns.size()); + logicalCost.liveOutCount = static_cast(stage.liveOuts.size()); + logicalCost.liveInBytes = stage.liveInBytes; + logicalCost.liveOutBytes = stage.liveOutBytes; + logicalCost.localSimtScopeCount = stage.localSimtScopeCount; + logicalCost.scopeInputTensorBytes = stage.scopeInputTensorBytes; + logicalCost.scopeOutputTensorBytes = stage.scopeOutputTensorBytes; + logicalCost.simtAnchorIndices = stage.simtAnchorIndices; + logicalCost.localSimtMaterializable = stage.localSimtMaterializable; + logicalCost.localSuperblockMaterializable = + stage.localSuperblockMaterializable; + logicalCost.legalSimtFactors = stage.legalSimtFactors; + logicalCost.localSimtFactors = stage.localSimtFactors; - for (const StageImplementation &implementation : implementations) { - if (!isDeclaredLegal(stage, implementation)) - return llvm::createStringError(std::errc::invalid_argument, - "Stage '%s' has an illegal candidate", - stage.id.c_str()); - StageResourceCycles resources = - mapWorkload(stage, - implementation.mode == StageMode::SIMD ? profile.simd - : profile.simt, - implementation.mode); - StageImplementationCost cost; - cost.implementation = implementation; - cost.resources = resources; - cost.totalCycles = applySuperBlock( - stage, resources, implementation, profile, - estimateStage(stage, profile, implementation.mode, resources)); - if (!cost.isValid()) - return llvm::createStringError(std::errc::invalid_argument, - "Stage '%s' produced an invalid cost", - stage.id.c_str()); - logicalCost.implementations.push_back(std::move(cost)); - } + llvm::SmallVector implementations; + if (stage.simdLegal) + implementations.push_back({StageMode::SIMD, 1, false}); + if (stage.simtLegal) + for (int64_t factor : stage.legalSimtFactors) + implementations.push_back({StageMode::SIMT, factor, false}); + if (stage.simtLegal && stage.localSimtMaterializable) + for (int64_t factor : stage.localSimtFactors) + implementations.push_back({StageMode::SIMT, factor, true}); - phaseCost.stages.push_back(logicalCost); - table.stages.push_back(std::move(logicalCost)); + for (const StageImplementation &implementation : implementations) { + if (!isDeclaredLegal(stage, implementation)) + return llvm::createStringError(std::errc::invalid_argument, + "Stage '%s' has an illegal candidate", + stage.id.c_str()); + StageResourceCycles resources = mapWorkload( + stage, + implementation.mode == StageMode::SIMD ? profile.simd : profile.simt, + implementation.mode); + StageImplementationCost cost; + cost.implementation = implementation; + cost.resources = resources; + cost.totalCycles = applySuperBlock( + stage, resources, implementation, profile, + estimateStage(stage, profile, implementation.mode, resources)); + if (!cost.isValid()) + return llvm::createStringError(std::errc::invalid_argument, + "Stage '%s' produced an invalid cost", + stage.id.c_str()); + logicalCost.implementations.push_back(std::move(cost)); } - table.phases.push_back(std::move(phaseCost)); + + table.stages.push_back(std::move(logicalCost)); } return table; } diff --git a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageRouteCostModel.cpp b/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageRouteCostModel.cpp index ac65f900a1..539a683252 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageRouteCostModel.cpp +++ b/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageRouteCostModel.cpp @@ -58,6 +58,24 @@ static double mixedEquivalentStageCost(const LogicalStageCost &stage, outputHandoffCycles; } +static double mixedExecutionMultiplicity(const LogicalStageCost &stage, + const StageImplementationCost &cost, + int64_t factor) { + if (factor <= 1 || cost.implementation.mode != StageMode::SIMD || + !stage.features.replicatedByLocalSuperBlock) + return 1.0; + // MaterializeSIMTScopeSuperBlock groups F logical programs in one physical + // V1-loop iteration. The selected SIMT scope executes them as one F-way + // VF, but every SIMD segment left in the logical body is cloned F times. + return static_cast(factor); +} + +static double mixedBaseStageCost(const LogicalStageCost &stage, + const StageImplementationCost &cost, + int64_t factor) { + return cost.totalCycles * mixedExecutionMultiplicity(stage, cost, factor); +} + /// AutoBlockify V1 is a route-conditional execution schedule. The analysis /// view contains its real dispatch/loop operations so pure-SIMT and Mixed can /// pay them, but an all-SIMD executable restores the original logical grid. @@ -189,6 +207,7 @@ llvm::json::Object StageModelFeatures::toJSON() const { result["has_reduction"] = hasReduction; result["has_dot"] = hasDot; result["has_conversion_pack"] = hasConversionPack; + result["replicated_by_local_superblock"] = replicatedByLocalSuperBlock; result["conditional_branch_count"] = conditionalBranchCount; result["divergent_branch_count"] = divergentBranchCount; result["loop_backedge_count"] = loopBackedgeCount; @@ -249,6 +268,10 @@ llvm::json::Object LogicalStageCost::toJSON() const { result["features"] = features.toJSON(); result["workload"] = workload.toJSON(); result["owned_operation_count"] = ownedOperationCount; + llvm::json::Array locations; + for (const std::string &location : sourceLocations) + locations.push_back(location); + result["source_locations"] = std::move(locations); result["live_in_count"] = liveInCount; result["live_out_count"] = liveOutCount; result["live_in_bytes"] = liveInBytes; @@ -261,6 +284,7 @@ llvm::json::Object LogicalStageCost::toJSON() const { anchorIndices.push_back(static_cast(index)); result["simt_anchor_indices"] = std::move(anchorIndices); result["local_simt_materializable"] = localSimtMaterializable; + result["local_superblock_materializable"] = localSuperblockMaterializable; llvm::json::Array legalFactors; for (int64_t factor : legalSimtFactors) legalFactors.push_back(factor); @@ -276,13 +300,6 @@ llvm::json::Object LogicalStageCost::toJSON() const { return result; } -llvm::json::Object LogicalPhaseCost::toJSON() const { - llvm::json::Array stageArray; - for (const LogicalStageCost &stage : stages) - stageArray.push_back(stage.toJSON()); - return llvm::json::Object{{"id", id}, {"stages", std::move(stageArray)}}; -} - bool StageTransitionCost::isValid() const { return std::isfinite(simdToSimtCycles) && std::isfinite(simtToSimdCycles) && simdToSimtCycles >= 0.0 && simtToSimdCycles >= 0.0 && @@ -333,25 +350,16 @@ llvm::json::Object StageRoutePlan::toJSON() const { stages.push_back(std::move(stage)); } result["stages"] = std::move(stages); - llvm::json::Array phases; - for (double cycles : logicalPhaseCycles) - phases.push_back(cycles); - result["logical_phase_system_cycles"] = std::move(phases); return result; } llvm::json::Object StageCostModelSummary::toJSON() const { llvm::json::Object result; result["applied"] = applied; - result["domain"] = domain; result["boundary_source"] = "operation_graph"; result["operation_ownership_complete"] = operationOwnershipComplete; result["modeled_operation_count"] = modeledOperationCount; result["profile_version"] = profileVersion; - llvm::json::Array phaseArray; - for (const LogicalPhaseCost &phase : phases) - phaseArray.push_back(phase.toJSON()); - result["logical_phases"] = std::move(phaseArray); llvm::json::Array stageArray; for (const LogicalStageCost &stage : stages) stageArray.push_back(stage.toJSON()); @@ -417,8 +425,9 @@ mlir::ascend::solveStageRoutes(const StageCostTable &costTable, findImplementation(stage, StageMode::SIMD, 1, false); const StageImplementationCost *simt = findImplementation(stage, StageMode::SIMT, factor, true); - const double simdCycles = - simd ? simd->totalCycles : std::numeric_limits::infinity(); + const double simdCycles = simd + ? mixedBaseStageCost(stage, *simd, factor) + : std::numeric_limits::infinity(); const double simtCycles = simt ? mixedEquivalentStageCost(stage, *simt, transition) : std::numeric_limits::infinity(); @@ -431,13 +440,72 @@ mlir::ascend::solveStageRoutes(const StageCostTable &costTable, if (kind != StageKernelRouteKind::Mixed) stageCycles = selected->totalCycles; - const double transitionCycles = stageCycles - selected->totalCycles; + const double transitionCycles = + stageCycles - mixedBaseStageCost(stage, *selected, factor); plan.implementations.push_back(selected->implementation); plan.entryTransitionCycles.push_back(transitionCycles); plan.logicalStageCycles.push_back(stageCycles); plan.totalCycles += stageCycles; } if (kind == StageKernelRouteKind::Mixed) { + if (factor > 1) { + // The current local-SuperBlock ABI widens one AutoBlockify V1 loop + // around exactly one selected local SIMT scope. Selecting two + // factor>1 scopes in the same loop would require preserving two + // independent warp/task mappings across the intervening SIMD + // segment, which the backend does not yet implement. Enumerate the + // one-SIMT-Stage routes instead of reporting an unmaterializable + // independent per-Stage minimum. + size_t bestSimtIndex = mixedChoices.size(); + double bestTotal = std::numeric_limits::infinity(); + for (size_t simtIndex = 0; simtIndex < mixedChoices.size(); + ++simtIndex) { + const MixedChoice &simtChoice = mixedChoices[simtIndex]; + if (!simtChoice.simt || + std::max( + 1, costTable.stages[simtIndex].localSimtScopeCount) != 1) + continue; + double candidateTotal = simtChoice.simtCycles; + bool candidateLegal = mixedChoices.size() > 1; + for (size_t index = 0; index < mixedChoices.size(); ++index) { + if (index == simtIndex) + continue; + if (!mixedChoices[index].simd) { + candidateLegal = false; + break; + } + candidateTotal += mixedChoices[index].simdCycles; + } + if (candidateLegal && candidateTotal < bestTotal) { + bestTotal = candidateTotal; + bestSimtIndex = simtIndex; + } + } + if (bestSimtIndex == mixedChoices.size()) { + StageRoutePlan invalid; + invalid.candidate = kind; + return invalid; + } + + plan.implementations.clear(); + plan.entryTransitionCycles.clear(); + plan.logicalStageCycles.clear(); + plan.totalCycles = 0.0; + for (size_t index = 0; index < mixedChoices.size(); ++index) { + const MixedChoice &choice = mixedChoices[index]; + const StageImplementationCost *selected = + index == bestSimtIndex ? choice.simt : choice.simd; + const double selectedCycles = + index == bestSimtIndex ? choice.simtCycles : choice.simdCycles; + plan.implementations.push_back(selected->implementation); + plan.entryTransitionCycles.push_back( + selectedCycles - + mixedBaseStageCost(costTable.stages[index], *selected, factor)); + plan.logicalStageCycles.push_back(selectedCycles); + plan.totalCycles += selectedCycles; + } + } + auto countMode = [&](StageMode mode) { return llvm::count_if(plan.implementations, [&](const auto &selected) { return selected.mode == mode; @@ -474,7 +542,8 @@ mlir::ascend::solveStageRoutes(const StageCostTable &costTable, plan.implementations[bestIndex] = replacement->implementation; plan.logicalStageCycles[bestIndex] = replacementCycles; plan.entryTransitionCycles[bestIndex] = - replacementCycles - replacement->totalCycles; + replacementCycles - mixedBaseStageCost(costTable.stages[bestIndex], + *replacement, factor); return true; }; @@ -519,11 +588,9 @@ mlir::ascend::solveStageRoutes(const StageCostTable &costTable, StageCostModelSummary result; result.applied = true; - result.domain = costTable.domain; result.operationOwnershipComplete = costTable.operationOwnershipComplete; result.modeledOperationCount = costTable.modeledOperationCount; result.profileVersion = costTable.profileVersion; - result.phases = costTable.phases; result.stages = costTable.stages; result.transition = transition; result.allSimd = buildPlan(StageKernelRouteKind::AllSIMD, 1); @@ -531,19 +598,5 @@ mlir::ascend::solveStageRoutes(const StageCostTable &costTable, result.mixed = bestFactoredPlan(StageKernelRouteKind::Mixed); removeAutoBlockifyCostFromAllSIMD(result.allSimd, costTable); - auto aggregatePhases = [&](StageRoutePlan &plan) { - if (!plan.legal) - return; - size_t stageIndex = 0; - for (const LogicalPhaseCost &phase : costTable.phases) { - double phaseCycles = 0.0; - for (size_t i = 0; i < phase.stages.size(); ++i) - phaseCycles += plan.logicalStageCycles[stageIndex++]; - plan.logicalPhaseCycles.push_back(phaseCycles); - } - }; - aggregatePhases(result.allSimd); - aggregatePhases(result.allSimt); - aggregatePhases(result.mixed); return result; } diff --git a/third_party/ascend/costmodel/lib/AscendModel/Transforms/SelectSimdSimtCostModel.cpp b/third_party/ascend/costmodel/lib/AscendModel/Transforms/SelectSimdSimtCostModel.cpp index cdf66e6b27..2b1a3163be 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/Transforms/SelectSimdSimtCostModel.cpp +++ b/third_party/ascend/costmodel/lib/AscendModel/Transforms/SelectSimdSimtCostModel.cpp @@ -195,31 +195,23 @@ struct SelectSimdSimtCostModelPass } SimdSimtCostReport report = std::move(*reportOr); - std::string recommended = - report.stageModel.applied - ? stringifySimdSimtCandidate(report.decision).str() - : kBackendDefault.str(); + std::string recommended = stringifySimdSimtCandidate(report.decision).str(); std::string effective = kBackendDefault.str(); std::string selectionSource = "backend_default"; std::string applicationReason; SmallVector mixedAnchors; SimtAnchorPlan selectedMixedAnchorPlan; int64_t selectedSuperblockFactor = 1; - if (report.stageModel.applied) { - if (report.decision == SimdSimtCandidateKind::AllSIMD) - selectedSuperblockFactor = - report.stageModel.allSimd.routeSuperblockFactor; - else if (report.decision == SimdSimtCandidateKind::AllSIMTOnly) - selectedSuperblockFactor = - report.stageModel.allSimt.routeSuperblockFactor; - else - selectedSuperblockFactor = - report.stageModel.mixed.routeSuperblockFactor; - } + if (report.decision == SimdSimtCandidateKind::AllSIMD) + selectedSuperblockFactor = + report.stageModel.allSimd.routeSuperblockFactor; + else if (report.decision == SimdSimtCandidateKind::AllSIMTOnly) + selectedSuperblockFactor = + report.stageModel.allSimt.routeSuperblockFactor; + else + selectedSuperblockFactor = report.stageModel.mixed.routeSuperblockFactor; - bool actionSupported = report.stageModel.applied; - if (!report.stageModel.applied) - applicationReason = "stage_model_not_applicable"; + bool actionSupported = true; bool hasExplicitScope = containsExplicitVectorScope(module); if (recommended == kMixedSimdSimt) { if (hasExplicitScope) { @@ -239,7 +231,7 @@ struct SelectSimdSimtCostModelPass } } // A factor>1 mixed route needs batching of the surrounding SIMD - // producer/consumer phases, not just a scope attribute. Keep the + // producer/consumer Stages, not just a scope attribute. Keep the // recommendation visible but do not apply it until ScopeSuperBlockPass // implements that exact materialization. if (selectedSuperblockFactor > 1 && @@ -273,7 +265,7 @@ struct SelectSimdSimtCostModelPass effective = recommended; selectionSource = "cpp_cost_model"; applicationReason = "minimum_cost_candidate"; - } else if (!autoMode && report.stageModel.applied) { + } else if (!autoMode) { applicationReason = "report_mode"; } else if (applicationReason.empty()) { applicationReason = "candidate_not_materializable"; @@ -285,20 +277,12 @@ struct SelectSimdSimtCostModelPass module->setAttr(kEffectiveExecutionAttr, builder.getStringAttr(effective)); module->setAttr(kSelectionSourceAttr, builder.getStringAttr(selectionSource)); - if (report.stageModel.applied) { - module->setAttr(kAllSimdScoreAttr, - builder.getF64FloatAttr(report.candidateCosts.allSimd)); - module->setAttr( - kAllSimtScoreAttr, - builder.getF64FloatAttr(report.candidateCosts.allSimtOnly)); - module->setAttr( - kMixedScoreAttr, - builder.getF64FloatAttr(report.candidateCosts.mixedSimdSimt)); - } else { - module->removeAttr(kAllSimdScoreAttr); - module->removeAttr(kAllSimtScoreAttr); - module->removeAttr(kMixedScoreAttr); - } + module->setAttr(kAllSimdScoreAttr, + builder.getF64FloatAttr(report.candidateCosts.allSimd)); + module->setAttr(kAllSimtScoreAttr, + builder.getF64FloatAttr(report.candidateCosts.allSimtOnly)); + module->setAttr(kMixedScoreAttr, builder.getF64FloatAttr( + report.candidateCosts.mixedSimdSimt)); module->setAttr(kSuperblockFactorAttr, builder.getI64IntegerAttr(selectedSuperblockFactor)); @@ -312,14 +296,6 @@ struct SelectSimdSimtCostModelPass } llvm::json::Object reportJSON = report.toJSON(); - if (!report.stageModel.applied) { - // No Legacy/aggregate fallback exists. Do not publish - // default-initialized candidate scores as if the Stage model had - // evaluated this kernel. - reportJSON.erase("candidate_costs"); - reportJSON["decision_kind"] = kBackendDefault; - reportJSON["selectable_candidates"] = llvm::json::Array(); - } reportJSON["mode"] = mode.getValue(); reportJSON["recommended_decision_kind"] = recommended; reportJSON["effective_decision_kind"] = effective; diff --git a/third_party/ascend/costmodel/profiles/microbench/ascend_davidv100_v1.json b/third_party/ascend/costmodel/profiles/microbench/ascend_davidv100_v1.json index cc1b22c691..5a8884fcfe 100644 --- a/third_party/ascend/costmodel/profiles/microbench/ascend_davidv100_v1.json +++ b/third_party/ascend/costmodel/profiles/microbench/ascend_davidv100_v1.json @@ -196,7 +196,7 @@ "confidence": "medium" }, "simt.setup.transition_harness_net.warps_1": { - "description": "Standalone serialized empty-VF net setup proxy from the branch-heavy transition harness for one launched warp. It is mode1 minus barrier-only mode6 and does not contain a SIMD phase, so it is neither a mixed transition nor a dependent handoff latency.", + "description": "Standalone serialized empty-VF net setup proxy from the branch-heavy transition harness for one launched warp. It is mode1 minus barrier-only mode6 and does not contain a SIMD Stage, so it is neither a mixed transition nor a dependent handoff latency.", "value": 182.0, "unit": "system_cycle", "cycle_domain": "SYS_CNT", @@ -206,7 +206,7 @@ "confidence": "low" }, "simt.setup.transition_harness_net.warps_2": { - "description": "Standalone serialized empty-VF net setup proxy from the branch-heavy transition harness for two launched warps. It is mode1 minus barrier-only mode6 and does not contain a SIMD phase, so it is neither a mixed transition nor a dependent handoff latency.", + "description": "Standalone serialized empty-VF net setup proxy from the branch-heavy transition harness for two launched warps. It is mode1 minus barrier-only mode6 and does not contain a SIMD Stage, so it is neither a mixed transition nor a dependent handoff latency.", "value": 182.0, "unit": "system_cycle", "cycle_domain": "SYS_CNT", @@ -216,7 +216,7 @@ "confidence": "low" }, "simt.setup.transition_harness_net.warps_4": { - "description": "Standalone serialized empty-VF net setup proxy from the branch-heavy transition harness for four launched warps. It is mode1 minus barrier-only mode6 and does not contain a SIMD phase, so it is neither a mixed transition nor a dependent handoff latency.", + "description": "Standalone serialized empty-VF net setup proxy from the branch-heavy transition harness for four launched warps. It is mode1 minus barrier-only mode6 and does not contain a SIMD Stage, so it is neither a mixed transition nor a dependent handoff latency.", "value": 182.0, "unit": "system_cycle", "cycle_domain": "SYS_CNT", @@ -226,7 +226,7 @@ "confidence": "low" }, "simt.setup.transition_harness_net.warps_8": { - "description": "Standalone serialized empty-VF net setup proxy from the branch-heavy transition harness for eight launched warps. It is mode1 minus barrier-only mode6 and does not contain a SIMD phase, so it is neither a mixed transition nor a dependent handoff latency.", + "description": "Standalone serialized empty-VF net setup proxy from the branch-heavy transition harness for eight launched warps. It is mode1 minus barrier-only mode6 and does not contain a SIMD Stage, so it is neither a mixed transition nor a dependent handoff latency.", "value": 182.0, "unit": "system_cycle", "cycle_domain": "SYS_CNT", @@ -236,7 +236,7 @@ "confidence": "low" }, "simt.setup.transition_harness_net.warps_16": { - "description": "Standalone serialized empty-VF net setup proxy from the branch-heavy transition harness for sixteen launched warps. It is mode1 minus barrier-only mode6 and does not contain a SIMD phase, so it is neither a mixed transition nor a dependent handoff latency.", + "description": "Standalone serialized empty-VF net setup proxy from the branch-heavy transition harness for sixteen launched warps. It is mode1 minus barrier-only mode6 and does not contain a SIMD Stage, so it is neither a mixed transition nor a dependent handoff latency.", "value": 182.0, "unit": "system_cycle", "cycle_domain": "SYS_CNT", @@ -246,7 +246,7 @@ "confidence": "low" }, "simt.setup.transition_harness_net.warps_32": { - "description": "Standalone serialized empty-VF net setup proxy from the branch-heavy transition harness for thirty-two launched warps. It is mode1 minus barrier-only mode6 and does not contain a SIMD phase, so it is neither a mixed transition nor a dependent handoff latency.", + "description": "Standalone serialized empty-VF net setup proxy from the branch-heavy transition harness for thirty-two launched warps. It is mode1 minus barrier-only mode6 and does not contain a SIMD Stage, so it is neither a mixed transition nor a dependent handoff latency.", "value": 223.0, "unit": "system_cycle", "cycle_domain": "SYS_CNT", diff --git a/third_party/ascend/costmodel/profiles/microbench/data_provider/camodel/README.md b/third_party/ascend/costmodel/profiles/microbench/data_provider/camodel/README.md index adc1cefe89..f97f5569ec 100644 --- a/third_party/ascend/costmodel/profiles/microbench/data_provider/camodel/README.md +++ b/third_party/ascend/costmodel/profiles/microbench/data_provider/camodel/README.md @@ -15,6 +15,7 @@ data_provider/ README.md # this workflow camodel_experiment_matrix.json # planned CAModel experiment coverage extract_camodel_system_cycle_profile.py + compare_stage_costs.py # Stage prediction / CaModel evidence join ``` ## 1. Build the probe @@ -74,7 +75,49 @@ source probe, host launcher, raw CAModel output, parser, and derived result. ## 3. Parse CAModel counts into SYS_CNT-domain rates -The parser expects a normalized JSON count file with this shape: +Run the following commands from `data_provider/`. The first parser accepts +either the `OPPROF_*` root directory or its `dump/` subdirectory. A supported +input contains primary SIMT issue/receive dumps such as: + +```text +OPPROF_YYYYMMDDHHMMSS_xxx/ + dump/ + core0.veccore0.rvec.simt.lsu.dump + core0.veccore0.rvec.simt.exu0.dump + core0.veccore0.rvec.simt.dvg0.dump + ... +``` + +First normalize the raw CAModel dumps into per-unit instruction counts: + +```bash +cd third_party/ascend/costmodel/profiles/microbench/data_provider + +python3 camodel/parse_camodel_counts.py \ + /path/to/OPPROF_YYYYMMDDHHMMSS_xxx \ + -o parsed_camodel_counts.json +``` + +Passing the `dump/` directory directly is equivalent: + +```bash +python3 camodel/parse_camodel_counts.py \ + /path/to/OPPROF_YYYYMMDDHHMMSS_xxx/dump \ + -o parsed_camodel_counts.json +``` + +Verify that at least one active unit was parsed before converting the rates: + +```bash +python3 -c \ + 'import json; d=json.load(open("parsed_camodel_counts.json")); print(d["per_unit"].keys())' +``` + +The command should print entries such as `core0.veccore0`. An empty mapping +means that the selected directory does not contain a supported primary SIMT +dump; do not continue with an empty file. + +The generated JSON count file has this shape: ```json { @@ -94,13 +137,25 @@ The parser expects a normalized JSON count file with this shape: } ``` -Then run: +Then convert the normalized counts from simulator cycles into SYS_CNT-domain +effective rates and save the result: ```bash python3 camodel/extract_camodel_system_cycle_profile.py parsed_camodel_counts.json \ --simulator-clock-mhz 1650.0 \ --sys-cnt-mhz 988.9 \ - --scope simt_memory + --scope simt_memory \ + > parsed_simt_memory.json +``` + +The complete data flow is: + +```text +OPPROF_* raw dumps + -> parse_camodel_counts.py + -> parsed_camodel_counts.json + -> extract_camodel_system_cycle_profile.py + -> parsed_simt_memory.json ``` The parser emits JSON like: @@ -151,3 +206,65 @@ Before treating a CAModel result as a data source: 5. `ascend_davidv100_v1.json` `source` points to all relevant artifacts. 6. The profile description does not claim hardware peak if the measurement is only workload-effective. + +## 6. Compare a selected route with Stage-level evidence + +`compare_stage_costs.py` joins the selected Stage implementations in a Route +Model report with an instruction-summary CSV from the same CaModel binary: + +```bash +python3 camodel/compare_stage_costs.py route_report.json \ + OPPROF_xxx/device0/core0.veccore1_instr_exe_xxx.csv \ + --route mixed --output stage_comparison.json +``` + +Without `--stage-pc-map`, the output deliberately keeps every Stage's CaModel +observation as `unobservable_without_stage_pc_map`; only the whole-kernel +resource-family totals are reported. This prevents whole-kernel residuals +from being presented as measured Stage cycles. + +For actual Stage attribution, the PC map must come from the compiler build of +the same binary and contain non-overlapping half-open ranges: + +```json +{ + "stages": [ + {"id": "head_index_mask", "pc_begin": "0x1000", "pc_end": "0x1100"}, + {"id": "diagonal_recurrence", "pc_begin": "0x1100", "pc_end": "0x1800"} + ] +} +``` + +Then pass `--stage-pc-map stage_pc_map.json`. CaModel simulator cycles remain +in their original clock domain; absolute comparison with Route Model system +cycles is invalid until the output also records the measured simulator/SYS_CNT +clock conversion. + +When the binary was built with `TRITON_DISABLE_LINE_INFO=false`, the simulator +also emits `*_code_exe.csv`. Pass it with +`--code-correlation-csv path/to/core0.veccore0_code_exe.csv`; the tool derives +source-line ownership from each Stage's `source_locations` report field. A +line shared by multiple Stages is reported in `ambiguous_source_lines` and is +not divided between them. This replaces the old workload-specific scripts +that hard-coded solve_tril line-number ranges. + +For finer evidence, the instruction CSV can be correlated directly with the +same debug binary. The load bias must be taken from that CaModel launch (it is +the runtime address corresponding to binary text address zero), not guessed +from a different run: + +```bash +python3 camodel/compare_stage_costs.py route_report.json \ + core0.veccore0_instr_exe.csv --route all_simt \ + --binary OPPROF_xxx/dump/aicore_binary.o \ + --load-bias 0x10d0d000 \ + --addr2line /path/to/llvm-addr2line \ + --output stage_instruction_comparison.json +``` + +The tool follows inline frames back to the Triton Python line, classifies the +instruction resource family, and attributes it only when that line and family +have exactly one predicted Stage owner. Ambiguous and unmatched evidence is +reported explicitly. Debug-line attribution is supplementary: lowering may +move a load to the line of a later dot/store, so compiler-emitted Stage PC +ranges remain the required evidence for complete Stage calibration. diff --git a/third_party/ascend/costmodel/profiles/microbench/data_provider/camodel/compare_stage_costs.py b/third_party/ascend/costmodel/profiles/microbench/data_provider/camodel/compare_stage_costs.py new file mode 100644 index 0000000000..346dcb70da --- /dev/null +++ b/third_party/ascend/costmodel/profiles/microbench/data_provider/camodel/compare_stage_costs.py @@ -0,0 +1,425 @@ +"""Compare selected StageCostModel terms with CaModel instruction evidence. + +The tool deliberately separates observation from attribution. An instruction +CSV can always provide a whole-kernel envelope and resource-family totals, but +it can only provide per-Stage measurements when a compiler-produced PC range +map is supplied. Missing attribution is reported as unobservable instead of +being filled with a residual or a workload-specific ratio. +""" + +import argparse +import csv +import json +import re +import subprocess +from collections import defaultdict +from pathlib import Path + +RESOURCE_KEYS = ( + "setup", + "load_per_iteration", + "store_per_iteration", + "compute_per_iteration", + "dot_per_iteration", + "scalar_per_iteration", + "predicate_per_iteration", + "shuffle_per_iteration", + "branch_control_per_iteration", + "loop_control_per_iteration", + "divergence_per_iteration", + "synchronization_per_iteration", + "spill_per_iteration", + "issue_per_iteration", + "critical_path_per_iteration", + "epilogue", +) + +FAMILY_RESOURCE_KEYS = { + "load": ("load_per_iteration", ), + "store": ("store_per_iteration", ), + "compute": ("compute_per_iteration", ), + "dot": ("dot_per_iteration", ), + "scalar": ("scalar_per_iteration", ), + "predicate": ("predicate_per_iteration", ), + "shuffle": ("shuffle_per_iteration", ), + "control": ("branch_control_per_iteration", "loop_control_per_iteration", "divergence_per_iteration"), + "synchronization": ("synchronization_per_iteration", ), + "spill": ("spill_per_iteration", ), +} + + +def parse_int(value): + return int(str(value), 0) + + +def classify_instruction(name, pipe): + upper_name = name.upper() + upper_pipe = pipe.upper() + if upper_name == "VF_SIMT": + return "envelope" + if "LD" in upper_pipe or upper_name.startswith(("SIMT_LD", "LD_")): + return "load" + if "ST" in upper_pipe or upper_name.startswith(("SIMT_ST", "ST_")): + return "store" + if "DOT" in upper_name or "MMAD" in upper_name or "CUBE" in upper_pipe: + return "dot" + if any(token in upper_name for token in ("SHFL", "SHUFFLE", "P2R", "R2P")): + return "shuffle" + if any(token in upper_name for token in ("SETP", "ISETP", "PRED")): + return "predicate" + if any(token in upper_name for token in ("BRANCH", "DVG", "LOOP")): + return "control" + if any(token in upper_name for token in ("SYNC", "SET_FLAG", "WAIT_FLAG")): + return "synchronization" + if upper_name.endswith(("STK", "LDK")): + return "spill" + if "SCALAR" in upper_pipe: + return "scalar" + return "compute" + + +def load_pc_map(path): + if path is None: + return [] + data = json.loads(path.read_text(encoding="utf-8")) + ranges = data.get("stages", data) + normalized = [] + for stage in ranges: + normalized.append({ + "id": stage["id"], + "begin": parse_int(stage["pc_begin"]), + "end": parse_int(stage["pc_end"]), + }) + return normalized + + +def stage_for_pc(pc, ranges): + matches = [entry["id"] for entry in ranges if entry["begin"] <= pc < entry["end"]] + if len(matches) > 1: + raise RuntimeError(f"overlapping Stage PC ranges for {pc:#x}: {matches}") + return matches[0] if matches else None + + +def read_camodel_csv(path, ranges): + aggregate = defaultdict(lambda: {"static_instructions": 0, "call_count": 0, "cycles": 0.0, "running_time_us": 0.0}) + per_stage = defaultdict(lambda: defaultdict(lambda: { + "static_instructions": 0, + "call_count": 0, + "cycles": 0.0, + "running_time_us": 0.0, + })) + unmapped = [] + with path.open(newline="", encoding="utf-8-sig") as stream: + for row in csv.DictReader(stream): + family = classify_instruction(row["instr"], row.get("pipe", "")) + values = { + "static_instructions": 1, + "call_count": int(row["call_count"]), + "cycles": float(row["cycles"]), + "running_time_us": float(row["running_time(us)"]), + } + for key, value in values.items(): + aggregate[family][key] += value + if not ranges or family == "envelope": + continue + pc = parse_int(row["addr"]) + stage_id = stage_for_pc(pc, ranges) + if stage_id is None: + unmapped.append(f"{pc:#x}") + continue + for key, value in values.items(): + per_stage[stage_id][family][key] += value + return aggregate, per_stage, sorted(set(unmapped)) + + +def addr2line(binary, addresses, load_bias, executable): + relative = [address - load_bias for address in addresses] + if any(address < 0 for address in relative): + raise RuntimeError("CaModel instruction address is below the supplied load bias") + command = [str(executable), "-a", "-i", "-e", str(binary), *(hex(address) for address in relative)] + result = subprocess.run(command, check=True, capture_output=True, text=True) + stacks = defaultdict(list) + current = None + for line in result.stdout.splitlines(): + if re.fullmatch(r"0x[0-9a-fA-F]+", line): + current = int(line, 16) + load_bias + elif current is not None: + stacks[current].append(line) + locations = {} + for address in addresses: + stack = stacks.get(address, []) + # Inline stacks often start in a CCE intrinsic or library helper. The + # Triton source frame is the evidence required by Stage provenance. + locations[address] = next((line for line in stack if ".py:" in line), stack[0] if stack else "??:0") + return locations + + +def family_is_predicted(stage, family): + keys = FAMILY_RESOURCE_KEYS.get(family, ()) + resources = stage["predicted_resource_system_cycles"] + return any(float(resources.get(key, 0.0)) > 0.0 for key in keys) + + +def read_instruction_source_attribution(path, predicted, binary, load_bias, executable): + with path.open(newline="", encoding="utf-8-sig") as stream: + rows = list(csv.DictReader(stream)) + addresses = sorted({parse_int(row["addr"]) for row in rows}) + locations = addr2line(binary, addresses, load_bias, executable) + source_owners = build_source_owners(predicted) + stage_by_id = {stage["id"]: stage for stage in predicted} + per_stage = defaultdict(lambda: defaultdict(lambda: { + "static_instructions": 0, + "call_count": 0, + "cycles": 0.0, + "running_time_us": 0.0, + })) + ambiguous = defaultdict(set) + unmatched = set() + for row in rows: + family = classify_instruction(row["instr"], row.get("pipe", "")) + if family == "envelope": + continue + key = source_line_key(locations[parse_int(row["addr"])]) + if key is None: + continue + candidates = { + stage_id + for stage_id in source_owners.get(key, set()) + if family_is_predicted(stage_by_id[stage_id], family) + } + evidence = f"{key}:{family}" + if not candidates: + unmatched.add(evidence) + continue + if len(candidates) != 1: + ambiguous[evidence].update(candidates) + continue + target = per_stage[next(iter(candidates))][family] + target["static_instructions"] += 1 + target["call_count"] += int(row["call_count"]) + target["cycles"] += float(row["cycles"]) + target["running_time_us"] += float(row["running_time(us)"]) + return per_stage, {key: sorted(value) for key, value in ambiguous.items()}, sorted(unmatched) + + +def source_line_key(value): + match = re.search(r'([^/\\"()]+\.py)"?:(\d+)', value) + if not match: + return None + return f"{match.group(1)}:{match.group(2)}" + + +def build_source_owners(predicted): + owners = defaultdict(set) + for stage in predicted: + for location in stage.get("source_locations", []): + key = source_line_key(location) + if key: + owners[key].add(stage["id"]) + return owners + + +def read_code_correlation(path, predicted): + owners = build_source_owners(predicted) + per_stage = defaultdict(lambda: { + "rows": 0, + "call_count": 0, + "cycles": 0.0, + "running_time_us": 0.0, + "source_lines": [], + }) + ambiguous = defaultdict(set) + unmatched = set() + with path.open(newline="", encoding="utf-8-sig") as stream: + for row in csv.DictReader(stream): + key = source_line_key(row.get("code", "")) + if key is None: + continue + stage_ids = owners.get(key, set()) + if not stage_ids: + unmatched.add(key) + continue + if len(stage_ids) != 1: + ambiguous[key].update(stage_ids) + continue + stage_id = next(iter(stage_ids)) + target = per_stage[stage_id] + target["rows"] += 1 + target["call_count"] += int(float(row.get("call_count") or 0)) + target["cycles"] += float(row.get("cycles") or 0) + target["running_time_us"] += float(row.get("running_time(us)") or 0) + target["source_lines"].append(key) + for value in per_stage.values(): + value["source_lines"] = sorted(set(value["source_lines"])) + return per_stage, {key: sorted(value) for key, value in ambiguous.items()}, sorted(unmatched) + + +def find_stage_model(report): + model = report.get("stage_model") or report.get("stage_cost_model") + if model is None: + raise RuntimeError("report does not contain stage_model") + return model + + +def flatten_stages(model): + if "logical_stages" not in model: + raise RuntimeError("stage_model does not contain logical_stages") + return model["logical_stages"] + + +def selected_route(model, route_name): + routes = model["routes"] + aliases = { + "all_simd": "all_simd", + "all_simt": "all_simt_only", + "mixed": "mixed_simd_simt", + } + route = routes[aliases[route_name]] + if not route.get("legal", False): + raise RuntimeError(f"requested route {route_name} is illegal") + return route + + +def implementation_for(stage, selection): + wanted = selection["implementation"] + for implementation in stage["implementations"]: + actual = implementation["implementation"] + if (actual["mode"] == wanted["mode"] + and int(actual.get("superblock_factor", 1)) == int(wanted.get("superblock_factor", 1)) + and bool(actual.get("local_scope", False)) == bool(wanted.get("local_scope", False))): + return implementation + raise RuntimeError(f"selected implementation missing from Stage {stage['id']}") + + +def predicted_stage_rows(model, route_name): + stages = flatten_stages(model) + route = selected_route(model, route_name) + selections = route["stages"] + if len(stages) != len(selections): + raise RuntimeError("route selection count does not match logical Stage count") + rows = [] + for stage, selection in zip(stages, selections): + implementation = implementation_for(stage, selection) + resources = implementation["resource_system_cycles"] + rows.append({ + "id": stage["id"], + "kind": stage.get("model"), + "source_locations": stage.get("source_locations", []), + "implementation": selection["implementation"], + "iteration_count": stage.get("iteration_count", 1), + "predicted_total_system_cycles": selection["logical_stage_system_cycles"], + "predicted_entry_transition_system_cycles": selection.get("entry_transition_system_cycles", 0.0), + "predicted_resource_system_cycles": {key: resources.get(key, 0.0) + for key in RESOURCE_KEYS}, + }) + return rows, route + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("report", type=Path) + parser.add_argument("instruction_csv", type=Path) + parser.add_argument("--route", choices=("all_simd", "all_simt", "mixed"), required=True) + parser.add_argument("--stage-pc-map", type=Path) + parser.add_argument("--code-correlation-csv", type=Path) + parser.add_argument("--binary", type=Path) + parser.add_argument("--load-bias", type=parse_int) + parser.add_argument("--addr2line", type=Path, default=Path("llvm-addr2line")) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + + report = json.loads(args.report.read_text(encoding="utf-8")) + model = find_stage_model(report) + predicted, route = predicted_stage_rows(model, args.route) + ranges = load_pc_map(args.stage_pc_map) + aggregate, per_stage, unmapped = read_camodel_csv(args.instruction_csv, ranges) + source_observations = {} + ambiguous_source_lines = {} + unmatched_source_lines = [] + if args.code_correlation_csv: + source_observations, ambiguous_source_lines, unmatched_source_lines = read_code_correlation( + args.code_correlation_csv, predicted) + instruction_source_observations = {} + ambiguous_instruction_sources = {} + unmatched_instruction_sources = [] + if args.binary: + if args.load_bias is None: + parser.error("--binary requires --load-bias") + (instruction_source_observations, ambiguous_instruction_sources, + unmatched_instruction_sources) = read_instruction_source_attribution(args.instruction_csv, predicted, + args.binary, args.load_bias, + args.addr2line) + + for row in predicted: + if ranges: + row["camodel_observation"] = { + "status": "observed" if row["id"] in per_stage else "no_matching_pc", + "resource_families": per_stage.get(row["id"], {}), + } + else: + row["camodel_observation"] = { + "status": "unobservable_without_stage_pc_map", + "resource_families": {}, + } + if args.code_correlation_csv: + row["camodel_source_observation"] = { + "status": "observed" if row["id"] in source_observations else "not_uniquely_attributed", + "totals": source_observations.get(row["id"], {}), + } + if args.binary: + row["camodel_instruction_source_observation"] = { + "status": "observed" if row["id"] in instruction_source_observations else "not_uniquely_attributed", + "resource_families": instruction_source_observations.get(row["id"], {}), + } + + output = { + "schema_version": + 1, + "unit_contract": { + "predicted": "system_cycle_selection_score", + "camodel_cycles": "simulator_cycle_aggregate_from_instruction_csv", + "camodel_running_time": "microseconds_reported_by_camodel", + "direct_numeric_comparison_valid": False, + }, + "route": + args.route, + "route_factor": + route.get("route_superblock_factor", 1), + "stage_assignment": + "pc_range" if ranges else "unobservable", + "stage_pc_map": + str(args.stage_pc_map) if args.stage_pc_map else None, + "code_correlation_csv": + str(args.code_correlation_csv) if args.code_correlation_csv else None, + "ambiguous_source_lines": + ambiguous_source_lines, + "unmatched_source_lines": + unmatched_source_lines, + "binary": + str(args.binary) if args.binary else None, + "load_bias": + args.load_bias, + "ambiguous_instruction_sources": + ambiguous_instruction_sources, + "unmatched_instruction_sources": + unmatched_instruction_sources, + "unmapped_instruction_pcs": + unmapped, + "camodel_kernel_resource_families": + aggregate, + "stages": + predicted, + "next_required_evidence": ([] if ranges else [ + "compiler-emitted non-overlapping Stage PC ranges, or uniquely attributable debug-line correlation, for the same binary", + "SYS_CNT conversion metadata before comparing absolute cycle values", + ]), + } + rendered = json.dumps(output, indent=2, sort_keys=True) + "\n" + if args.output: + args.output.write_text(rendered, encoding="utf-8", newline="\n") + else: + print(rendered, end="") + + +if __name__ == "__main__": + main() diff --git a/third_party/ascend/costmodel/profiles/microbench/data_provider/camodel/parse_camodel_counts.py b/third_party/ascend/costmodel/profiles/microbench/data_provider/camodel/parse_camodel_counts.py new file mode 100755 index 0000000000..d19a276b28 --- /dev/null +++ b/third_party/ascend/costmodel/profiles/microbench/data_provider/camodel/parse_camodel_counts.py @@ -0,0 +1,228 @@ +#!/usr/bin/env python3 +"""Parse raw CAModel output into normalized per-unit instruction counts.""" + +import argparse +import csv +import json +import re +from collections import Counter, defaultdict +from pathlib import Path + +SIMT_OP_RE = re.compile(r"\[(\d+)\].*?\b(SIMT_[A-Z0-9_]+)\b") +CORE_RE = re.compile(r"(core\d+)\.(veccore\d+)\.rvec\.simt\.(.+)\.dump$") + +PRIMARY_SUFFIXES = ( + "lsu", + "dvg0", + "dvg1", + "dvg2", + "dvg3", + "exu0", + "exu1", + "exu2", + "exu3", +) + +OP_GROUPS = { + "memory": ("SIMT_LDG", "SIMT_STG", "SIMT_LDS", "SIMT_STS"), + "shuffle": ("SIMT_SHFL", ), + "predicate": ("SIMT_ISETP", "SIMT_ISETP_I", "SIMT_PLOP3"), + "control": ("SIMT_BRANCH", "SIMT_END"), + "float_alu": ("SIMT_FADD", "SIMT_FMUL", "SIMT_FMNMX", "SIMT_FMNMX_I"), + "int_alu": ("SIMT_IADD", "SIMT_IADD_I", "SIMT_IADD_X", "SIMT_IADD_X_I", "SIMT_IMUL", "SIMT_SHFI", "SIMT_LOP3"), + "move": ("SIMT_MOV", ), +} + + +def _as_float(value): + try: + return float(value) + except (TypeError, ValueError): + return 0.0 + + +def _iter_primary_simt_dumps(dump_dir): + for path in sorted(Path(dump_dir).glob("core*.veccore*.rvec.simt.*.dump")): + match = CORE_RE.search(path.name) + if not match: + continue + suffix = match.group(3) + if suffix in PRIMARY_SUFFIXES: + yield path, match.group(1), match.group(2), suffix + + +def _scan_simt_dump(path): + counts = Counter() + first_ts = {} + last_ts = {} + examples = {} + with path.open(errors="ignore") as file: + for line in file: + match = SIMT_OP_RE.search(line) + if not match: + continue + timestamp = int(match.group(1)) + op = match.group(2) + counts[op] += 1 + first_ts.setdefault(op, timestamp) + last_ts[op] = timestamp + examples.setdefault(op, line.strip()[:180]) + return counts, first_ts, last_ts, examples + + +def _merge_counter(dst, src): + for key, value in src.items(): + dst[key] += value + + +def _group_counts(op_counts): + groups = Counter() + matched = set() + for group, prefixes in OP_GROUPS.items(): + for op, count in op_counts.items(): + if any(op == prefix or op.startswith(prefix + "_") for prefix in prefixes): + groups[group] += count + matched.add(op) + for op, count in op_counts.items(): + if op not in matched: + groups["other"] += count + return groups + + +def _span_from_first_last(first, last): + if not first or not last: + return None + return { + "first": min(first.values()), + "last": max(last.values()), + "delta": max(last.values()) - min(first.values()), + } + + +def _parse_instr_exe(root): + result = { + "files": [], + "pipe_cycles": Counter(), + "instr_cycles": Counter(), + "instr_calls": Counter(), + } + for path in sorted(Path(root).glob("**/*instr_exe*.csv")): + with path.open(newline="", errors="ignore") as file: + rows = list(csv.DictReader(file)) + pipe_cycles = Counter() + instr_cycles = Counter() + instr_calls = Counter() + for row in rows: + instr = row.get("instr") or row.get("Instr") or row.get("instruction") or "" + pipe = row.get("pipe") or row.get("Pipe") or "" + cycles = _as_float(row.get("cycles") or row.get("Cycles") or row.get("cycle")) + calls = _as_float(row.get("call_count") or row.get("Call Count") or row.get("count") or 1) + pipe_cycles[pipe] += cycles + instr_cycles[instr] += cycles + instr_calls[instr] += calls + result["files"].append({ + "path": str(path), + "rows": len(rows), + "pipe_cycles": dict(pipe_cycles.most_common()), + "instr_cycles": dict(instr_cycles.most_common(50)), + }) + _merge_counter(result["pipe_cycles"], pipe_cycles) + _merge_counter(result["instr_cycles"], instr_cycles) + _merge_counter(result["instr_calls"], instr_calls) + result["pipe_cycles"] = dict(result["pipe_cycles"].most_common()) + result["instr_cycles"] = dict(result["instr_cycles"].most_common(80)) + result["instr_calls"] = dict(result["instr_calls"].most_common(80)) + return result + + +def _make_seed(op_counts, group_counts, span): + total_ops = sum(op_counts.values()) + duration = span["delta"] if span else 0 + seed = { + "duration_cycles": duration, + "total_primary_simt_ops": total_ops, + "naive_cycles_per_primary_op": duration / total_ops if total_ops else None, + "groups": {}, + } + for group, count in group_counts.items(): + seed["groups"][group] = { + "ops": count, + "ops_per_cycle": count / duration if duration else None, + "naive_cycles_per_op": duration / count if count else None, + } + return seed + + +def extract(root): + root = Path(root) + dump_dir = root / "dump" if (root / "dump").is_dir() else root + aggregate_counts = Counter() + aggregate_first = {} + aggregate_last = {} + aggregate_examples = {} + per_unit = defaultdict(lambda: {"counts": Counter(), "first_ts": {}, "last_ts": {}, "files": []}) + + for path, core, veccore, _suffix in _iter_primary_simt_dumps(dump_dir): + counts, first_ts, last_ts, examples = _scan_simt_dump(path) + unit_key = f"{core}.{veccore}" + per_unit[unit_key]["files"].append(path.name) + _merge_counter(per_unit[unit_key]["counts"], counts) + _merge_counter(aggregate_counts, counts) + for op, timestamp in first_ts.items(): + per_unit[unit_key]["first_ts"].setdefault(op, timestamp) + aggregate_first.setdefault(op, timestamp) + aggregate_examples.setdefault(op, examples.get(op, "")) + for op, timestamp in last_ts.items(): + per_unit[unit_key]["last_ts"][op] = max(per_unit[unit_key]["last_ts"].get(op, timestamp), timestamp) + aggregate_last[op] = max(aggregate_last.get(op, timestamp), timestamp) + + per_unit_output = {} + for unit, data in sorted(per_unit.items()): + counts = data["counts"] + groups = _group_counts(counts) + span = _span_from_first_last(data["first_ts"], data["last_ts"]) + per_unit_output[unit] = { + "files": data["files"], + "op_counts": dict(counts.most_common()), + "group_counts": dict(groups.most_common()), + "span": span, + "seed": _make_seed(counts, groups, span), + } + + aggregate_groups = _group_counts(aggregate_counts) + aggregate_span = _span_from_first_last(aggregate_first, aggregate_last) + return { + "source": str(root), + "dump_dir": str(dump_dir), + "primary_suffixes": list(PRIMARY_SUFFIXES), + "aggregate": { + "op_counts": dict(aggregate_counts.most_common()), + "group_counts": dict(aggregate_groups.most_common()), + "span": aggregate_span, + "examples": {op: aggregate_examples[op] + for op in aggregate_counts + if op in aggregate_examples}, + "seed": _make_seed(aggregate_counts, aggregate_groups, aggregate_span), + }, + "per_unit": per_unit_output, + "instr_exe": _parse_instr_exe(root), + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("root", type=Path, help="OPPROF directory or CAModel dump directory") + parser.add_argument("-o", "--output", type=Path, help="output normalized JSON path") + args = parser.parse_args() + + result = extract(args.root) + if not result["per_unit"]: + parser.error(f"no supported primary SIMT dumps found under {result['dump_dir']}") + payload = json.dumps(result, indent=2, sort_keys=True) + if args.output: + args.output.write_text(payload + "\n", encoding="utf-8") + print(payload) + + +if __name__ == "__main__": + main() diff --git a/third_party/ascend/costmodel/profiles/simd_simt/david_v100_simd_simt_v1.json b/third_party/ascend/costmodel/profiles/simd_simt/david_v100_simd_simt_v1.json index 67155babd3..ca6e30556a 100644 --- a/third_party/ascend/costmodel/profiles/simd_simt/david_v100_simd_simt_v1.json +++ b/third_party/ascend/costmodel/profiles/simd_simt/david_v100_simd_simt_v1.json @@ -287,7 +287,7 @@ "useful_factor_limit": 4, "persistent_state_pressure_free_factor": 2, "persistent_state_bytes_per_system_cycle": 8.0, - "description": "Independent logical programs may continue hiding latency through F4. Loop-carried live state is pressure-free only through F2 and is charged above that point. Keeping the two limits separate lets straight-line FBGEMM-like work benefit from F4 without making recurrence-heavy solve_tril pressure-free." + "description": "Independent logical programs may continue hiding latency through F4. The 8 B/system-cycle persistent-state rate is shared by the generic SIMD loop-carried-state term and by the extra pressure of replicating that state in SIMT SuperBlock; state replication is pressure-free only through F2. The rule consumes SSA live-out bytes and contains no workload-name check." }, "scope_handoff": { "fixed_directional_system_cycles": 668.0, @@ -325,7 +325,7 @@ "measurement": "simt.setup.transition_harness_net.warps_32" }, "confidence": "low", - "note": "These values are mode1(empty SIMT plus tail barrier) minus mode6(barrier only) in a branch-heavy harness. They contain no SIMD phase and are used only as conservative setup fallbacks. Directional SIMD-to-SIMT, SIMT-to-SIMD, and dependent UB hand-off costs remain unmeasured." + "note": "These values are mode1(empty SIMT plus tail barrier) minus mode6(barrier only) in a branch-heavy harness. They contain no SIMD Stage and are used only as conservative setup fallbacks. Directional SIMD-to-SIMT, SIMT-to-SIMD, and dependent UB hand-off costs remain unmeasured." } } } diff --git a/third_party/ascend/costmodel/profiles/simd_simt/simd_simt_profile_schema.json b/third_party/ascend/costmodel/profiles/simd_simt/simd_simt_profile_schema.json index 6652e397d6..bb6d739e23 100644 --- a/third_party/ascend/costmodel/profiles/simd_simt/simd_simt_profile_schema.json +++ b/third_party/ascend/costmodel/profiles/simd_simt/simd_simt_profile_schema.json @@ -1049,7 +1049,7 @@ }, "mixed_setup_fallbacks": { "type": "object", - "description": "Map from warp count to low-confidence mixed setup fallback. The referenced probes contain no SIMD phase and are not directional transition measurements.", + "description": "Map from warp count to low-confidence mixed setup fallback. The referenced probes contain no SIMD Stage and are not directional transition measurements.", "required": [ "1", "2", diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/simt_scope.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/simt_scope.mlir index e86daa3bc8..f93e7a75e3 100644 --- a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/simt_scope.mlir +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/simt_scope.mlir @@ -1,16 +1,12 @@ // RUN: triton-opt --triton-to-linalg="named-ops=True" --split-input-file %s | FileCheck %s -// A canonical SIMT scope must select both the mixed BiShengIR pipeline and -// the SIMT-aware runtime launch path. +// A canonical SIMT scope must select the mixed BiShengIR pipeline and the +// SIMT-aware runtime launch path. The scope is an input contract and may be +// consumed by outlining/lowering, so this test guards the surviving function +// contract rather than requiring the source scope operation in output IR. // CHECK-LABEL: func.func @simt_scope // CHECK-SAME: parallel_mode = "mix_simd_simt" -// CHECK: scope.scope -// CHECK: } { -// CHECK-SAME: hivm.func_core_type = #hivm.func_core_type -// CHECK-SAME: hivm.vf_mode = #hivm.vf_mode -// CHECK-SAME: no_inline -// CHECK-SAME: outline -// CHECK-SAME: vector_mode = "simt" +// CHECK-NOT: parallel_mode = "mix_simd" tt.func public @simt_scope(%arg0: !tt.ptr) { %zero = arith.constant dense<0.000000e+00> : tensor<16xf32> %ptrs = tt.splat %arg0 : !tt.ptr -> tensor<16x!tt.ptr> diff --git a/third_party/ascend/unittest/costmodel_ut/PassesTest.cpp b/third_party/ascend/unittest/costmodel_ut/PassesTest.cpp index 51b63bb1ae..36fa59a8b3 100644 --- a/third_party/ascend/unittest/costmodel_ut/PassesTest.cpp +++ b/third_party/ascend/unittest/costmodel_ut/PassesTest.cpp @@ -217,7 +217,6 @@ module { if (!features) FAIL() << llvm::toString(features.takeError()); - EXPECT_EQ(features->loadedIndexDependentMemoryOps, 1); EXPECT_EQ(features->simtAnchors.count, 1); } @@ -518,7 +517,7 @@ TEST(CostModelPassesTest, PerfReportPassAcceptsEstimatedPipeline) { createPipelineAnalysisPass(), createPerfReportPass())); } -TEST(CostModelPassesTest, SimdSimtReturnsBackendDefaultOutsideStageDomain) { +TEST(CostModelPassesTest, SimdSimtScoresGenericSemanticStages) { auto configureOptions = [](SelectSimdSimtCostModelPassOptions &options, llvm::StringRef mode) { options.mode = mode.str(); @@ -548,19 +547,19 @@ TEST(CostModelPassesTest, SimdSimtReturnsBackendDefaultOutsideStageDomain) { ASSERT_TRUE(autoEffective); ASSERT_TRUE(autoRecommended); ASSERT_TRUE(autoReport); - EXPECT_EQ(autoEffective.getValue(), "backend_default"); - EXPECT_EQ(autoRecommended.getValue(), "backend_default"); - EXPECT_FALSE((*autoModule)->hasAttr("ascend.simt_costmodel.all_simd_score")); + EXPECT_NE(autoEffective.getValue(), "backend_default"); + EXPECT_EQ(autoEffective.getValue(), autoRecommended.getValue()); + EXPECT_TRUE((*autoModule)->hasAttr("ascend.simt_costmodel.all_simd_score")); auto autoJSON = llvm::json::parse(autoReport.getValue()); ASSERT_TRUE(static_cast(autoJSON)); auto *autoObject = autoJSON->getAsObject(); ASSERT_NE(autoObject, nullptr); auto autoDecision = autoObject->getString("decision_kind"); ASSERT_TRUE(autoDecision); - EXPECT_EQ(*autoDecision, "backend_default"); + EXPECT_NE(*autoDecision, "backend_default"); auto autoReason = autoObject->getString("application_reason"); ASSERT_TRUE(autoReason); - EXPECT_EQ(*autoReason, "stage_model_not_applicable"); + EXPECT_EQ(*autoReason, "minimum_cost_candidate"); mlir::MLIRContext reportContext; auto reportModule = parseModule(reportContext, kOutOfSimdSimtCoverageModule); @@ -579,18 +578,17 @@ TEST(CostModelPassesTest, SimdSimtReturnsBackendDefaultOutsideStageDomain) { ASSERT_TRUE(reportEffective); ASSERT_TRUE(reportJSONAttr); EXPECT_EQ(reportEffective.getValue(), "backend_default"); - EXPECT_FALSE( - (*reportModule)->hasAttr("ascend.simt_costmodel.all_simd_score")); + EXPECT_TRUE((*reportModule)->hasAttr("ascend.simt_costmodel.all_simd_score")); auto reportJSON = llvm::json::parse(reportJSONAttr.getValue()); ASSERT_TRUE(static_cast(reportJSON)); auto *reportObject = reportJSON->getAsObject(); ASSERT_NE(reportObject, nullptr); auto reportDecision = reportObject->getString("decision_kind"); ASSERT_TRUE(reportDecision); - EXPECT_EQ(*reportDecision, "backend_default"); + EXPECT_NE(*reportDecision, "backend_default"); auto reportReason = reportObject->getString("application_reason"); ASSERT_TRUE(reportReason); - EXPECT_EQ(*reportReason, "stage_model_not_applicable"); + EXPECT_EQ(*reportReason, "report_mode"); } TEST(CostModelPassesTest, SimdSimtSelectionUsesExternalAnalysisIR) { diff --git a/third_party/ascend/unittest/costmodel_ut/SimdSimtCostModelTest.cpp b/third_party/ascend/unittest/costmodel_ut/SimdSimtCostModelTest.cpp index 539997ebbc..a5e6d33c9f 100644 --- a/third_party/ascend/unittest/costmodel_ut/SimdSimtCostModelTest.cpp +++ b/third_party/ascend/unittest/costmodel_ut/SimdSimtCostModelTest.cpp @@ -10,8 +10,6 @@ #include using mlir::ascend::HardwareProfile; -using mlir::ascend::LogicalPhase; -using mlir::ascend::LogicalPhaseCost; using mlir::ascend::LogicalStage; using mlir::ascend::LogicalStageCost; using mlir::ascend::SimdSimtFeatureSummary; @@ -36,10 +34,6 @@ namespace { SimdSimtFeatureSummary triangularBt16StageFeatures() { SimdSimtFeatureSummary f; - f.reduceOps = 1; - f.loadOps = 1; - f.storeOps = 1; - f.staticLoopTripCountMax = 14; f.simtAnchors.count = 1; TriangularSolveFacts triangular; triangular.blockRows = 16; @@ -115,12 +109,8 @@ logicalStage(llvm::StringRef id, StageCostModelKind kind, llvm::Expected evaluateOneStage(LogicalStage stage, HardwareProfile profile = hardwareProfile()) { - LogicalPhase phase; - phase.id = "phase"; - phase.stages.push_back(std::move(stage)); StagePartition partition; - partition.domain = "unit_test"; - partition.phases.push_back(std::move(phase)); + partition.stages.push_back(std::move(stage)); return StageCostEvaluator().evaluate(partition, profile); } @@ -142,9 +132,6 @@ TEST(SimdSimtCostModelTest, ScopeSuperBlockLegalityRequiresBackendAndResourceMaximum) { auto makePartition = [](int64_t independentGroups) { StagePartition partition; - partition.domain = "unit_test"; - mlir::ascend::LogicalPhase phase; - phase.id = "phase"; LogicalStage stage = logicalStage("payload", StageCostModelKind::LoopCarriedRecurrence, StageScheduleKind::LoopCarriedSerial, /*iterations=*/16); @@ -152,23 +139,22 @@ TEST(SimdSimtCostModelTest, stage.features.hasLoopCarriedDataDependency = true; stage.features.parallelRecurrenceGroupCount = independentGroups; stage.localSimtMaterializable = true; + stage.localSuperblockMaterializable = true; stage.localSimtFactors = {1}; - phase.stages.push_back(std::move(stage)); - partition.phases.push_back(std::move(phase)); + partition.stages.push_back(std::move(stage)); return partition; }; StagePartition f1Only = makePartition(/*independentGroups=*/4); if (llvm::Error error = StageModeLegalityAnalysis().analyze(f1Only, 4, false)) FAIL() << llvm::toString(std::move(error)); - EXPECT_EQ(f1Only.phases[0].stages[0].localSimtFactors, - (std::vector{1})); + EXPECT_EQ(f1Only.stages[0].localSimtFactors, (std::vector{1})); StagePartition scopeSuperblock = makePartition(/*independentGroups=*/4); if (llvm::Error error = StageModeLegalityAnalysis().analyze(scopeSuperblock, 4, true)) FAIL() << llvm::toString(std::move(error)); - EXPECT_EQ(scopeSuperblock.phases[0].stages[0].localSimtFactors, + EXPECT_EQ(scopeSuperblock.stages[0].localSimtFactors, (std::vector{1, 2, 4})); // ABI-v2 creates an F1 V1 scheduling loop and refines only the selected @@ -178,11 +164,9 @@ TEST(SimdSimtCostModelTest, if (llvm::Error error = StageModeLegalityAnalysis().analyze(mixedOnly, 1, true)) FAIL() << llvm::toString(std::move(error)); - EXPECT_EQ(mixedOnly.phases[0].stages[0].legalSimtFactors, - (std::vector{1})); - EXPECT_EQ(mixedOnly.phases[0].stages[0].localSimtFactors, - (std::vector{1})); - auto mixedOnlyCosts = evaluateOneStage(mixedOnly.phases[0].stages[0]); + EXPECT_EQ(mixedOnly.stages[0].legalSimtFactors, (std::vector{1})); + EXPECT_EQ(mixedOnly.stages[0].localSimtFactors, (std::vector{1})); + auto mixedOnlyCosts = evaluateOneStage(mixedOnly.stages[0]); if (!mixedOnlyCosts) FAIL() << llvm::toString(mixedOnlyCosts.takeError()); ASSERT_EQ(mixedOnlyCosts->stages[0].implementations.size(), 3u); @@ -195,35 +179,31 @@ TEST(SimdSimtCostModelTest, if (llvm::Error error = StageModeLegalityAnalysis().analyze(oneWorkGroup, 4, true)) FAIL() << llvm::toString(std::move(error)); - EXPECT_EQ(oneWorkGroup.phases[0].stages[0].localSimtFactors, + EXPECT_EQ(oneWorkGroup.stages[0].localSimtFactors, (std::vector{1, 2, 4})); } TEST(SimdSimtCostModelTest, LocalScopeFactorsHonorKernelResourceMaximum) { StagePartition partition; - partition.domain = "local_factor_limit"; - LogicalPhase phase; - phase.id = "gather"; LogicalStage stage; stage.id = "indirect_tile_gather"; stage.costModelKind = StageCostModelKind::IndirectGatherMemory; stage.scheduleKind = StageScheduleKind::PartiallyDependent; stage.iterationCount = 1; stage.localSimtMaterializable = true; - phase.stages.push_back(std::move(stage)); - partition.phases.push_back(std::move(phase)); + stage.localSuperblockMaterializable = true; + partition.stages.push_back(std::move(stage)); ASSERT_FALSE(StageModeLegalityAnalysis().analyze( partition, /*maximumSuperblockFactor=*/2, /*scopeSuperblockMaterializable=*/true)); - const LogicalStage &result = partition.phases.front().stages.front(); + const LogicalStage &result = partition.stages.front(); EXPECT_EQ(result.legalSimtFactors, (std::vector{1, 2})); EXPECT_EQ(result.localSimtFactors, (std::vector{1, 2})); } TEST(SimdSimtCostModelTest, KernelMixedRouteComesFromAdjacentStageModes) { StageCostTable table; - table.domain = "unit_test"; table.profileVersion = "unit-test-profile-v1"; auto addStage = [&](llvm::StringRef id, double simd, double simt) { mlir::ascend::LogicalStageCost stage; @@ -244,10 +224,6 @@ TEST(SimdSimtCostModelTest, KernelMixedRouteComesFromAdjacentStageModes) { addStage("head", 10.0, 20.0); addStage("payload", 100.0, 50.0); addStage("store", 30.0, 45.0); - mlir::ascend::LogicalPhaseCost phase; - phase.id = "kernel"; - phase.stages = table.stages; - table.phases.push_back(std::move(phase)); StageTransitionCost transition; transition.simdToSimtCycles = 5.0; @@ -268,7 +244,6 @@ TEST(SimdSimtCostModelTest, KernelMixedRouteComesFromAdjacentStageModes) { TEST(SimdSimtCostModelTest, MixedScopePaysExactBidirectionalUbHandoffCost) { StageCostTable table; - table.domain = "scope_handoff"; table.profileVersion = "unit-test-profile-v1"; auto makeCost = [&](StageMode mode, double cycles, bool localScope = false) { mlir::ascend::StageImplementationCost cost; @@ -293,10 +268,6 @@ TEST(SimdSimtCostModelTest, MixedScopePaysExactBidirectionalUbHandoffCost) { mlir::ascend::LogicalStageCost tail = head; tail.id = "tail"; table.stages = {head, payload, tail}; - mlir::ascend::LogicalPhaseCost phase; - phase.id = "phase"; - phase.stages = table.stages; - table.phases.push_back(std::move(phase)); StageTransitionCost transition; transition.simdUbLoadBytesPerCycle = 512.0; @@ -321,7 +292,6 @@ TEST(SimdSimtCostModelTest, MixedScopePaysExactBidirectionalUbHandoffCost) { TEST(SimdSimtCostModelTest, MixedScopeSuperBlockAmortizesOnlyFixedTransitions) { StageCostTable table; - table.domain = "scope_superblock_handoff"; table.profileVersion = "unit-test-profile-v1"; auto makeCost = [&](StageMode mode, int64_t factor, double cycles, bool localScope = false) { @@ -352,10 +322,6 @@ TEST(SimdSimtCostModelTest, MixedScopeSuperBlockAmortizesOnlyFixedTransitions) { LogicalStageCost tail = head; tail.id = "tail"; table.stages = {head, payload, tail}; - LogicalPhaseCost phase; - phase.id = "phase"; - phase.stages = table.stages; - table.phases.push_back(std::move(phase)); StageTransitionCost transition; transition.simdToSimtCycles = 40.0; @@ -490,6 +456,28 @@ TEST(SimdSimtCostModelTest, RecurrenceAccumulatesCriticalPathAndTraffic) { EXPECT_GT(table->stages[0].implementations[0].resources.criticalPath, 0.0); } +TEST(SimdSimtCostModelTest, SimdRecurrenceChargesPersistentLiveStateBytes) { + LogicalStage baseline = logicalStage( + "baseline_recurrence", StageCostModelKind::LoopCarriedRecurrence, + StageScheduleKind::LoopCarriedSerial); + baseline.features.hasLoop = true; + baseline.features.hasLoopCarriedDataDependency = true; + LogicalStage withState = baseline; + withState.id = "stateful_recurrence"; + withState.liveOutBytes = 800; + + auto baselineTable = evaluateOneStage(baseline); + auto stateTable = evaluateOneStage(withState); + if (!baselineTable) + FAIL() << llvm::toString(baselineTable.takeError()); + if (!stateTable) + FAIL() << llvm::toString(stateTable.takeError()); + const double baselineSimd = + baselineTable->stages[0].implementations[0].totalCycles; + const double stateSimd = stateTable->stages[0].implementations[0].totalCycles; + EXPECT_DOUBLE_EQ(stateSimd - baselineSimd, 800.0 / 8.0); +} + TEST(SimdSimtCostModelTest, SimtRecurrenceInterleavesIndependentGroupsButKeepsIssueFloor) { LogicalStage serial = logicalStage("serial_recurrence", @@ -613,7 +601,6 @@ TEST(SimdSimtCostModelTest, IndirectMemoryUsesDependencyProfile) { TEST(SimdSimtCostModelTest, MixedRouteRejectsUnmaterializableSimtStage) { StageCostTable table; - table.domain = "unit_test"; table.profileVersion = "unit-test-profile-v1"; auto makeCost = [&](StageMode mode, double cycles, bool localScope = false) { mlir::ascend::StageImplementationCost cost; @@ -632,10 +619,6 @@ TEST(SimdSimtCostModelTest, MixedRouteRejectsUnmaterializableSimtStage) { payload.implementations = {makeCost(StageMode::SIMD, 100.0), makeCost(StageMode::SIMT, 1.0)}; table.stages = {head, payload}; - mlir::ascend::LogicalPhaseCost phase; - phase.id = "phase"; - phase.stages = table.stages; - table.phases.push_back(std::move(phase)); auto routes = solveStageRoutes(table, StageTransitionCost{}); if (!routes) FAIL() << llvm::toString(routes.takeError()); @@ -646,7 +629,6 @@ TEST(SimdSimtCostModelTest, MixedRouteRejectsUnmaterializableSimtStage) { TEST(SimdSimtCostModelTest, MixedRouteReportsCheapestConstrainedRouteWhenLocalScopeLoses) { StageCostTable table; - table.domain = "constrained_mixed"; table.profileVersion = "unit-test-profile-v1"; auto makeCost = [&](StageMode mode, double cycles, bool localScope = false) { mlir::ascend::StageImplementationCost cost; @@ -667,10 +649,6 @@ TEST(SimdSimtCostModelTest, dot.implementations = {makeCost(StageMode::SIMD, 40.0), makeCost(StageMode::SIMT, 90.0)}; table.stages = {gather, dot}; - mlir::ascend::LogicalPhaseCost phase; - phase.id = "gather_dot_min"; - phase.stages = table.stages; - table.phases.push_back(std::move(phase)); StageTransitionCost transition; transition.simdToSimtCycles = 10.0; @@ -688,7 +666,6 @@ TEST(SimdSimtCostModelTest, TEST(SimdSimtCostModelTest, AllSimdDoesNotPayRouteConditionalAutoBlockify) { StageCostTable table; - table.domain = "auto_blockify_route_conditional"; table.profileVersion = "unit-test-profile-v1"; auto makeCost = [&](StageMode mode, double cycles) { mlir::ascend::StageImplementationCost cost; @@ -708,10 +685,6 @@ TEST(SimdSimtCostModelTest, AllSimdDoesNotPayRouteConditionalAutoBlockify) { payload.implementations = {makeCost(StageMode::SIMD, 100.0), makeCost(StageMode::SIMT, 80.0)}; table.stages = {dispatch, payload}; - mlir::ascend::LogicalPhaseCost phase; - phase.id = "phase"; - phase.stages = table.stages; - table.phases.push_back(std::move(phase)); auto routes = solveStageRoutes(table, StageTransitionCost{}); if (!routes) @@ -727,7 +700,6 @@ TEST(SimdSimtCostModelTest, AllSimdDoesNotPayRouteConditionalAutoBlockify) { TEST(SimdSimtCostModelTest, MixedRouteChargesEveryMaterializedScope) { StageCostTable table; - table.domain = "scope_count"; table.profileVersion = "unit-test-profile-v1"; auto makeCost = [&](StageMode mode, double cycles, bool localScope = false) { mlir::ascend::StageImplementationCost cost; @@ -750,10 +722,6 @@ TEST(SimdSimtCostModelTest, MixedRouteChargesEveryMaterializedScope) { mlir::ascend::LogicalStageCost tail = head; tail.id = "tail"; table.stages = {head, gather, tail}; - mlir::ascend::LogicalPhaseCost phase; - phase.id = "phase"; - phase.stages = table.stages; - table.phases.push_back(std::move(phase)); StageTransitionCost transition; transition.simdToSimtCycles = 10.0; @@ -794,7 +762,6 @@ TEST(SimdSimtCostModelTest, SuperBlockLatencyHidingStopsAtUsefulFactorLimit) { TEST(SimdSimtCostModelTest, PureSimtRouteUsesOneUniformSuperBlockFactor) { StageCostTable table; - table.domain = "uniform_superblock"; table.profileVersion = "unit-test-profile-v1"; auto makeCost = [&](int64_t factor, double cycles) { mlir::ascend::StageImplementationCost cost; @@ -811,10 +778,6 @@ TEST(SimdSimtCostModelTest, PureSimtRouteUsesOneUniformSuperBlockFactor) { second.implementations = {makeCost(1, 5.0), makeCost(2, 4.0), makeCost(4, 1.0)}; table.stages = {first, second}; - mlir::ascend::LogicalPhaseCost phase; - phase.id = "phase"; - phase.stages = table.stages; - table.phases.push_back(std::move(phase)); auto routes = solveStageRoutes(table, StageTransitionCost{}); if (!routes) @@ -829,7 +792,6 @@ TEST(SimdSimtCostModelTest, PureSimtRouteUsesOneUniformSuperBlockFactor) { TEST(SimdSimtCostModelTest, MixedScopeSuperBlockUsesSelectedFactorCost) { StageCostTable table; - table.domain = "mixed_scope_superblock"; table.profileVersion = "unit-test-profile-v1"; auto makeCost = [&](StageMode mode, int64_t factor, double cycles, bool localScope = false) { @@ -862,10 +824,6 @@ TEST(SimdSimtCostModelTest, MixedScopeSuperBlockUsesSelectedFactorCost) { payload.localSimtMaterializable = true; payload.localSimtFactors = {1, 2, 4}; table.stages = {prefix, payload}; - mlir::ascend::LogicalPhaseCost phase; - phase.id = "phase"; - phase.stages = table.stages; - table.phases.push_back(std::move(phase)); auto routes = solveStageRoutes(table, StageTransitionCost{}); if (!routes) @@ -875,6 +833,51 @@ TEST(SimdSimtCostModelTest, MixedScopeSuperBlockUsesSelectedFactorCost) { EXPECT_DOUBLE_EQ(routes->mixed.totalCycles, 5.5); } +TEST(SimdSimtCostModelTest, + FactoredMixedRouteUsesOneBackendMaterializableLocalScope) { + StageCostTable table; + table.profileVersion = "unit-test-profile-v1"; + auto makeCost = [&](StageMode mode, int64_t factor, double cycles, + bool localScope = false) { + StageImplementationCost cost; + cost.implementation = {mode, factor, localScope}; + cost.totalCycles = cycles; + return cost; + }; + + LogicalStageCost first; + first.id = "first_local_candidate"; + first.features.replicatedByLocalSuperBlock = true; + first.localSimtMaterializable = true; + first.localSimtScopeCount = 1; + first.implementations = {makeCost(StageMode::SIMD, 1, 100.0), + makeCost(StageMode::SIMT, 4, 1.0, true)}; + LogicalStageCost second = first; + second.id = "second_local_candidate"; + LogicalStageCost tail; + tail.id = "simd_tail"; + tail.implementations = {makeCost(StageMode::SIMD, 1, 1.0)}; + table.stages = {first, second, tail}; + + auto routes = solveStageRoutes(table, StageTransitionCost{}); + if (!routes) + FAIL() << llvm::toString(routes.takeError()); + ASSERT_TRUE(routes->mixed.legal); + EXPECT_EQ(routes->mixed.routeSuperblockFactor, 4); + EXPECT_EQ(llvm::count_if( + routes->mixed.implementations, + [](const mlir::ascend::StageImplementation &implementation) { + return implementation.localScope; + }), + 1); + // One scope is SIMT (1 cycle); the other Stage remains SIMD and is cloned + // once per grouped logical program (100 * F4); the outside tail is not. + EXPECT_DOUBLE_EQ(routes->mixed.totalCycles, 402.0); + EXPECT_DOUBLE_EQ(routes->mixed.entryTransitionCycles[0], 0.0); + EXPECT_DOUBLE_EQ(routes->mixed.entryTransitionCycles[1], 0.0); + EXPECT_DOUBLE_EQ(routes->mixed.entryTransitionCycles[2], 0.0); +} + TEST(SimdSimtCostModelTest, OperationGraphBoundaryOwnsEveryRootAndDerivesLiveValues) { mlir::MLIRContext context; @@ -918,34 +921,25 @@ TEST(SimdSimtCostModelTest, mlir::ascend::SimtAnchorPlan anchorPlan; anchorPlan.anchors.push_back(std::move(anchor)); - SimdSimtFeatureSummary features = triangularBt16StageFeatures(); - auto phasePlan = mlir::ascend::PhaseBoundaryAnalysis().analyze( - *module, anchorPlan, features, StagePartitionerOptions{}); - if (!phasePlan) - FAIL() << llvm::toString(phasePlan.takeError()); - ASSERT_TRUE(*phasePlan); - EXPECT_EQ((*phasePlan)->rootOperations.size(), - (*phasePlan)->rootPhaseIds.size()); - EXPECT_EQ((*phasePlan)->rootPhaseIds, - std::vector({"head", "head", "head", "diagonal_load", - "diagonal_inverse", "merge_store"})); - - auto result = StagePartitioner().partition(*module, anchorPlan, features, + auto structure = + mlir::ascend::ProgramStructureAnalysis().analyze(*module, anchorPlan); + if (!structure) + FAIL() << llvm::toString(structure.takeError()); + EXPECT_EQ(structure->rootOperations.size(), 6u); + + auto result = StagePartitioner().partition(*module, anchorPlan, StagePartitionerOptions{}); if (!result) FAIL() << llvm::toString(result.takeError()); - ASSERT_TRUE(*result); - const StagePartition &partition = **result; + const StagePartition &partition = *result; EXPECT_TRUE(partition.operationOwnershipComplete); int64_t ownedRootCount = 0; const LogicalStage *recurrenceStage = nullptr; - for (const LogicalPhase &phase : partition.phases) { - for (const LogicalStage &stage : phase.stages) { - ownedRootCount += static_cast(stage.operations.size()); - if (stage.id == "diagonal_inverse_recurrence") - recurrenceStage = &stage; - } + for (const LogicalStage &stage : partition.stages) { + ownedRootCount += static_cast(stage.operations.size()); + if (llvm::is_contained(stage.operations, recurrence)) + recurrenceStage = &stage; } EXPECT_EQ(ownedRootCount, partition.modeledOperationCount); ASSERT_NE(recurrenceStage, nullptr); @@ -963,7 +957,53 @@ TEST(SimdSimtCostModelTest, } TEST(SimdSimtCostModelTest, - CompoundScopeOrderIsNormalizedBeforePhasePartitioning) { + SameStatementSupportOperationsJoinTheDominantResourceStage) { + mlir::MLIRContext context; + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.allowUnregisteredDialects(); + auto module = mlir::parseSourceString(R"mlir( + module { + func.func @kernel(%pointer: i64) { + %index = arith.constant 0 : i64 + %mask = arith.cmpi eq, %index, %index : i64 + %value = "tt.load"(%pointer, %mask) : (i64, i1) -> f32 + %tail = arith.constant 1 : i64 + return + } + } + )mlir", + &context); + ASSERT_TRUE(module); + + llvm::SmallVector roots; + module->walk([&](mlir::func::FuncOp function) { + for (mlir::Operation &operation : function.getBody().front()) + if (!operation.hasTrait()) + roots.push_back(&operation); + }); + ASSERT_EQ(roots.size(), 4u); + auto statement = mlir::FileLineColLoc::get(&context, "kernel.py", 10, 1); + for (mlir::Operation *operation : llvm::ArrayRef(roots).take_front(3)) + operation->setLoc(statement); + roots.back()->setLoc(mlir::FileLineColLoc::get(&context, "kernel.py", 11, 1)); + + mlir::ascend::ProgramStructure structure; + structure.rootOperations.assign(roots.begin(), roots.end()); + auto result = mlir::ascend::StageBoundaryAnalysis().analyze( + structure, mlir::ascend::SimtAnchorPlan{}); + if (!result) + FAIL() << llvm::toString(result.takeError()); + ASSERT_EQ(result->stages.size(), 2u); + EXPECT_EQ(result->stages.front().operations.size(), 3u); + EXPECT_EQ(result->stages.front().costModelKind, + StageCostModelKind::ContinuousTileMemory); + ASSERT_EQ(result->stages.back().operations.size(), 1u); + EXPECT_EQ(result->stages.back().operations.front(), roots.back()); +} + +TEST(SimdSimtCostModelTest, + CompoundScopeOrderIsNormalizedBeforeStagePartitioning) { mlir::MLIRContext context; context.getOrLoadDialect(); context.getOrLoadDialect(); @@ -1015,32 +1055,30 @@ TEST(SimdSimtCostModelTest, mlir::ascend::SimtAnchorPlan anchorPlan; anchorPlan.anchors.push_back(std::move(anchor)); - auto phasePlan = mlir::ascend::PhaseBoundaryAnalysis().analyze( - *module, anchorPlan, triangularBt16StageFeatures(), - StagePartitionerOptions{}); - if (!phasePlan) - FAIL() << llvm::toString(phasePlan.takeError()); - ASSERT_TRUE(*phasePlan); - EXPECT_EQ((*phasePlan)->rootPhaseIds, - std::vector({"head", "head", "head", "diagonal_load", - "diagonal_inverse", "diagonal_inverse", - "merge_store"})); + auto structure = + mlir::ascend::ProgramStructureAnalysis().analyze(*module, anchorPlan); + if (!structure) + FAIL() << llvm::toString(structure.takeError()); + auto setupPosition = llvm::find(structure->rootOperations, setup); + auto recurrencePosition = llvm::find(structure->rootOperations, recurrence); + ASSERT_NE(setupPosition, structure->rootOperations.end()); + ASSERT_NE(recurrencePosition, structure->rootOperations.end()); + EXPECT_EQ(recurrencePosition - setupPosition, 1); auto partition = StagePartitioner().partition(*module, anchorPlan, - triangularBt16StageFeatures(), StagePartitionerOptions{}); if (!partition) FAIL() << llvm::toString(partition.takeError()); - ASSERT_TRUE(*partition); const LogicalStage *loadStage = nullptr; const LogicalStage *recurrenceStage = nullptr; - for (const LogicalPhase &phase : (**partition).phases) - for (const LogicalStage &stage : phase.stages) { - if (stage.id == "load_diagonal_tiles") - loadStage = &stage; - if (stage.id == "diagonal_inverse_recurrence") - recurrenceStage = &stage; - } + for (const LogicalStage &stage : partition->stages) { + if (llvm::any_of(stage.operations, [&](mlir::Operation *operation) { + return operation->getName().getStringRef() == "tt.load"; + })) + loadStage = &stage; + if (llvm::is_contained(stage.operations, recurrence)) + recurrenceStage = &stage; + } ASSERT_NE(loadStage, nullptr); ASSERT_NE(recurrenceStage, nullptr); EXPECT_EQ(loadStage->operations.size(), 1u); @@ -1048,6 +1086,158 @@ TEST(SimdSimtCostModelTest, EXPECT_EQ(recurrenceStage->simtAnchorIndices, std::vector({0})); } +TEST(SimdSimtCostModelTest, + NestedLocalScopeDoesNotAdvertiseUnsupportedSuperBlockFactors) { + mlir::MLIRContext context; + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.allowUnregisteredDialects(); + auto module = mlir::parseSourceString(R"mlir( + module { + func.func @kernel(%pointer: i64, %condition: i1) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c4 = arith.constant 4 : index + scf.for %i = %c0 to %c4 step %c1 { + scf.if %condition { + %value = "tt.load"(%pointer) : (i64) -> tensor<16xf32> + } + } + return + } + } + )mlir", + &context); + ASSERT_TRUE(module); + + mlir::scf::ForOp v1Loop; + mlir::Operation *nestedLoad = nullptr; + module->walk([&](mlir::Operation *operation) { + if (auto loop = llvm::dyn_cast(operation)) + v1Loop = loop; + if (operation->getName().getStringRef() == "tt.load") + nestedLoad = operation; + }); + ASSERT_TRUE(v1Loop); + ASSERT_NE(nestedLoad, nullptr); + v1Loop->setAttr("ta.auto_blockify_v1.loop", mlir::UnitAttr::get(&context)); + + mlir::ascend::SimtAnchorDescriptor anchor; + anchor.operation = nestedLoad; + anchor.scopeOperations.push_back(nestedLoad); + anchor.scopeInsertionPoint = nestedLoad; + anchor.kind = mlir::ascend::SimtAnchorKind::DirectGather; + anchor.lowerability.mixed = true; + anchor.materializable = true; + mlir::ascend::SimtAnchorPlan anchorPlan; + anchorPlan.anchors.push_back(std::move(anchor)); + + StagePartitionerOptions options; + options.maximumSuperblockFactor = 4; + options.scopeSuperblockMaterializable = true; + auto result = StagePartitioner().partition(*module, anchorPlan, options); + if (!result) + FAIL() << llvm::toString(result.takeError()); + + const LogicalStage *nestedStage = nullptr; + for (const LogicalStage &stage : result->stages) + if (!stage.simtAnchorIndices.empty()) + nestedStage = &stage; + ASSERT_NE(nestedStage, nullptr); + EXPECT_TRUE(nestedStage->localSimtMaterializable); + EXPECT_FALSE(nestedStage->localSuperblockMaterializable); + EXPECT_FALSE(nestedStage->operations.empty()); + EXPECT_EQ(nestedStage->localSimtFactors, (std::vector{1})); + + auto costs = StageCostEvaluator().evaluate(*result, hardwareProfile()); + if (!costs) + FAIL() << llvm::toString(costs.takeError()); + auto nestedCost = llvm::find_if(costs->stages, [&](const auto &stage) { + return stage.id == nestedStage->id; + }); + ASSERT_NE(nestedCost, costs->stages.end()); + EXPECT_FALSE(nestedCost->sourceLocations.empty()); +} + +TEST(SimdSimtCostModelTest, GenericSemanticStagesDoNotRequireAWorkloadDomain) { + mlir::MLIRContext context; + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.allowUnregisteredDialects(); + auto module = mlir::parseSourceString(R"mlir( + module { + func.func @unrelated_kernel(%pointer: i64) { + %zero = arith.constant dense<0.0> : tensor<16xf32> + %loaded = "tt.load"(%pointer) : (i64) -> tensor<16xf32> + %mask = arith.cmpf ogt, %loaded, %zero : tensor<16xf32> + "tt.store"(%pointer, %loaded) : (i64, tensor<16xf32>) -> () + return + } + } + )mlir", + &context); + ASSERT_TRUE(module); + + mlir::ascend::SimtAnchorPlan anchorPlan; + auto partition = StagePartitioner().partition(*module, anchorPlan, + StagePartitionerOptions{}); + if (!partition) + FAIL() << llvm::toString(partition.takeError()); + + ASSERT_TRUE(partition->operationOwnershipComplete); + EXPECT_EQ(partition->modeledOperationCount, 4); + ASSERT_EQ(partition->stages.size(), 4u); + EXPECT_EQ(partition->stages[0].costModelKind, + StageCostModelKind::ScalarIssue); + EXPECT_EQ(partition->stages[1].costModelKind, + StageCostModelKind::ContinuousTileMemory); + EXPECT_EQ(partition->stages[2].costModelKind, + StageCostModelKind::PredicateMask); + EXPECT_EQ(partition->stages[3].costModelKind, + StageCostModelKind::ContinuousTileStore); +} + +TEST(SimdSimtCostModelTest, AdjacentStructuredLoopsRemainSerialStages) { + mlir::MLIRContext context; + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.allowUnregisteredDialects(); + auto module = mlir::parseSourceString(R"mlir( + module { + func.func @two_serial_loops(%pointer: i64) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c4 = arith.constant 4 : index + scf.for %i = %c0 to %c4 step %c1 { + %first = "tt.load"(%pointer) : (i64) -> f32 + } + scf.for %i = %c0 to %c4 step %c1 { + %second = "tt.load"(%pointer) : (i64) -> f32 + } + return + } + } + )mlir", + &context); + ASSERT_TRUE(module); + + mlir::ascend::SimtAnchorPlan anchorPlan; + auto partition = StagePartitioner().partition(*module, anchorPlan, + StagePartitionerOptions{}); + if (!partition) + FAIL() << llvm::toString(partition.takeError()); + + llvm::SmallVector loopStages; + for (const LogicalStage &stage : partition->stages) + if (stage.costModelKind == StageCostModelKind::IndependentPipelinedLoop) + loopStages.push_back(&stage); + ASSERT_EQ(loopStages.size(), 2u); + EXPECT_NE(loopStages[0]->operations.front(), + loopStages[1]->operations.front()); +} + TEST(SimdSimtCostModelTest, LocalScopeReturningPointerTensorIsRejectedBeforeScoring) { mlir::MLIRContext context; @@ -1089,22 +1279,17 @@ TEST(SimdSimtCostModelTest, mlir::ascend::SimtAnchorPlan anchorPlan; anchorPlan.anchors.push_back(std::move(anchor)); - mlir::ascend::PhaseBoundaryPlan phasePlan{ - mlir::ascend::PhaseBoundaryDomain::LoadedIndexRowwiseReduction, - "loaded_index_rowwise_reduction", std::nullopt}; - phasePlan.rootOperations.assign(roots.begin(), roots.end()); - phasePlan.rootPhaseIds = {"row_dispatch", "row_load", "row_load", - "row_reduction", "convert_store"}; - auto result = mlir::ascend::StageBoundaryAnalysis().analyze( - phasePlan, SimdSimtFeatureSummary{}, &anchorPlan); + mlir::ascend::ProgramStructure structure; + structure.rootOperations.assign(roots.begin(), roots.end()); + auto result = + mlir::ascend::StageBoundaryAnalysis().analyze(structure, anchorPlan); if (!result) FAIL() << llvm::toString(result.takeError()); const LogicalStage *gather = nullptr; - for (const LogicalPhase &phase : result->phases) - for (const LogicalStage &stage : phase.stages) - if (stage.id == "indirect_row_gather") - gather = &stage; + for (const LogicalStage &stage : result->stages) + if (llvm::is_contained(stage.operations, roots[1])) + gather = &stage; ASSERT_NE(gather, nullptr); EXPECT_FALSE(gather->localSimtMaterializable); EXPECT_TRUE(gather->localSimtFactors.empty()); @@ -1142,21 +1327,18 @@ TEST(SimdSimtCostModelTest, PointerInductionLoopIsNotADataRecurrence) { StagePartition partition; partition.operationOwnershipComplete = true; - LogicalPhase phase; - phase.id = "convert_store"; LogicalStage stage = logicalStage("pointer_loop", StageCostModelKind::ConversionPack, StageScheduleKind::IndependentPipelined, 8); stage.operations.push_back(loop); - phase.stages.push_back(std::move(stage)); - partition.phases.push_back(std::move(phase)); + partition.stages.push_back(std::move(stage)); if (llvm::Error error = StageFeatureAnalysis().analyze(partition)) FAIL() << llvm::toString(std::move(error)); if (llvm::Error error = mlir::ascend::StageKindClassifier().analyze(partition, 8192)) FAIL() << llvm::toString(std::move(error)); - const LogicalStage &classified = partition.phases.front().stages.front(); + const LogicalStage &classified = partition.stages.front(); EXPECT_TRUE(classified.features.hasLoop); EXPECT_TRUE(classified.features.hasPointerInduction); EXPECT_FALSE(classified.features.hasLoopCarriedDataDependency); @@ -1167,15 +1349,12 @@ TEST(SimdSimtCostModelTest, PointerInductionLoopIsNotADataRecurrence) { TEST(SimdSimtCostModelTest, IncompatibleDominantStructuresRequireStageSplit) { StagePartition partition; partition.operationOwnershipComplete = true; - LogicalPhase phase; - phase.id = "compound"; LogicalStage stage = logicalStage("gather_dot", StageCostModelKind::TinyCubeRoofline, StageScheduleKind::PartiallyDependent, 1); stage.features.hasDot = true; stage.features.hasIndirectMemory = true; - phase.stages.push_back(std::move(stage)); - partition.phases.push_back(std::move(phase)); + partition.stages.push_back(std::move(stage)); llvm::Error error = mlir::ascend::StageKindClassifier().analyze(partition, 16384); diff --git a/third_party/ascend/unittest/costmodel_ut/test_compiler_costmodel_contract.py b/third_party/ascend/unittest/costmodel_ut/test_compiler_costmodel_contract.py index 2523d708e6..250d73d812 100644 --- a/third_party/ascend/unittest/costmodel_ut/test_compiler_costmodel_contract.py +++ b/third_party/ascend/unittest/costmodel_ut/test_compiler_costmodel_contract.py @@ -230,6 +230,20 @@ def test_mixed_compile_keeps_delayed_cross_core_gss_enabled(self): source = inspect.getsource(cmplr.linalg_to_bin_enable_npu_compile_910_95) self.assertNotIn("--enable-hivm-delayed-cross-core-gss=false", source) + def test_all_bishengir_entries_share_debug_info_option(self): + cmplr, _dump_mgr, _GPUTarget = self._load_compiler_module() + + options = [] + cmplr._append_debug_info_option(options) + self.assertEqual(options, ["--enable-debug-info=true"]) + source = inspect.getsource(cmplr.ttir_to_npubin) + self.assertIn("_append_debug_info_option(_compile_option_list)", source) + + cmplr._is_debug_line_info_disabled = lambda: True + options = [] + cmplr._append_debug_info_option(options) + self.assertEqual(options, []) + if __name__ == "__main__": unittest.main() diff --git a/third_party/ascend/unittest/pytest_ut/test_simd_simt_costmodel_cases.py b/third_party/ascend/unittest/pytest_ut/test_simd_simt_costmodel_cases.py index 4d4cd107fe..395bd4e98f 100644 --- a/third_party/ascend/unittest/pytest_ut/test_simd_simt_costmodel_cases.py +++ b/third_party/ascend/unittest/pytest_ut/test_simd_simt_costmodel_cases.py @@ -162,7 +162,7 @@ def launch(): expected = torch.matmul(a[:, indices].float(), b[indices, :].float()) torch.testing.assert_close(output, expected, rtol=1e-2, atol=1e-2) report = _load_route_report(report_path, "all_simt_only") - assert report["features"]["dot_ops"] == 1 + assert any(stage["features"]["has_dot"] for stage in report["stage_model"]["logical_stages"]) _assert_performance("gather_dot_min", launch, tmp_path / "gather_profile", 5.478, tolerance=1.35) @@ -448,10 +448,26 @@ def launch(): launch() gathered = input_tensor[token_indices].float() * scores[token_indices, expert_indices].float()[:, None] row_max = torch.clamp(torch.amax(torch.abs(gathered), dim=1), min=1.0e-12) - expected_scale = row_max / 448.0 - expected = torch.clamp(gathered / expected_scale[:, None], -448.0, 448.0).to(torch.float8_e4m3fn) + # Match the kernel's floating-point operation order. Although + # ``x / (row_max / 448)`` is algebraically equivalent to + # ``x * (448 / row_max)``, their FP32 rounding differs at FP8 bin + # boundaries and can select adjacent values whose spacing is 32. + quant_scale = 448.0 / row_max + expected_scale = 1.0 / quant_scale + expected = torch.clamp(gathered * quant_scale[:, None], -448.0, 448.0).to(torch.float8_e4m3fn) torch.testing.assert_close(output_scale, expected_scale, rtol=2e-3, atol=2e-3) - torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=16) + output_f32 = output.float() + expected_f32 = expected.float() + difference = torch.abs(output_f32 - expected_f32) + # The device conversion and torch's reference conversion may choose + # adjacent FP8 values for an exact rounding tie. Check one E4M3 ULP at + # each expected value instead of using a fixed tolerance: E4M3 spacing is + # 16 around 128 but 32 around 256, while subnormals have spacing 2^-9. + magnitude = torch.abs(expected_f32) + normal_ulp = torch.pow(2.0, torch.floor(torch.log2(torch.clamp(magnitude, min=2**-6))) - 3) + fp8_ulp = torch.where(magnitude < 2**-6, torch.full_like(magnitude, 2**-9), normal_ulp) + assert torch.all(difference <= fp8_ulp), (f"FBGEMM FP8 output exceeds one ULP: max_abs={difference.max().item()}, " + f"max_ulp_error={(difference / fp8_ulp).max().item()}") layout_merge_disabled = os.getenv("TRITON_TEST_DISABLE_TTIR_LAYOUT_MERGE") == "1" expected_route = "all_simd" if layout_merge_disabled else "all_simt_only" report = _load_route_report(report_path, expected_route) From b3ee63dac63fddbe7ea6491c92f672af17ec3b59 Mon Sep 17 00:00:00 2001 From: yangkaixin Date: Thu, 3 Sep 2026 17:22:25 +0800 Subject: [PATCH 5/5] 1. asert to trunk_coalescing list, 2. add scan model --- third_party/ascend/backend/compiler.py | 6 ++-- third_party/ascend/backend/driver.py | 2 +- .../AscendModel/RouteModel/StageCostModels.h | 2 ++ .../RouteModel/StageRouteCostModel.h | 1 + .../Analysis/SimtAnchorAnalysis.cpp | 5 ---- .../AscendModel/Analysis/StagePartitioner.cpp | 10 ++++++- .../RouteModel/SimdSimtCostModel.cpp | 7 +++-- .../RouteModel/StageCostModels.cpp | 16 ++++++++++- .../RouteModel/StageRouteCostModel.cpp | 1 + .../simd_simt/david_v100_simd_simt_v1.json | 10 ++++++- .../simd_simt/simd_simt_profile_schema.json | 18 ++++++++++++ .../TritonToLinalg/TileChunkCoalescing.cpp | 28 ++++++++++++++++++- .../costmodel_ut/SimdSimtCostModelTest.cpp | 23 +++++++++++++++ .../pytest_ut/test_auto_blockify_v1_policy.py | 4 +-- .../pytest_ut/test_launcher_export_api.py | 2 +- .../test_simd_simt_costmodel_cases.py | 9 ++++-- 16 files changed, 124 insertions(+), 20 deletions(-) diff --git a/third_party/ascend/backend/compiler.py b/third_party/ascend/backend/compiler.py index 3c0faabdd4..7dee32c948 100644 --- a/third_party/ascend/backend/compiler.py +++ b/third_party/ascend/backend/compiler.py @@ -288,9 +288,9 @@ def _publish_route_transform_capability(metadata, opt) -> str: capability = { "schema_version": 1, "layout_merge_applied": bool(metadata.get("ttir_layout_merge_applied", False)), - "row_coalescing_applied": coalesce_factor > 1, - "row_coalescing_factor": coalesce_factor, - "row_coalescing_axis": coalesce_axis, + "layout_coalescing_applied": coalesce_factor > 1, + "layout_coalescing_factor": coalesce_factor, + "layout_coalescing_axis": coalesce_axis, "auto_blockify_v1_requested": bool(metadata.get("auto_blockify_v1_requested", False)), "auto_blockify_v1_materializable": v1_materializable, "auto_blockify_v1_disable_reasons": sorted(set(disable_reasons)), diff --git a/third_party/ascend/backend/driver.py b/third_party/ascend/backend/driver.py index e1e1afdd63..a4b7ab7099 100644 --- a/third_party/ascend/backend/driver.py +++ b/third_party/ascend/backend/driver.py @@ -708,7 +708,7 @@ def _format_of_msprof_task_type_ratio(bs_task_type, mix_mode): # along that axis, so the host shrinks the matching grid dim by H here (the # equivalent of what bishengir AutoBlockify used to do via hacc.coalesce_factor; # bishengir no longer touches it). Tile/strided merges require exact - # divisibility; independent-row coalescing carries a runtime tail predicate + # divisibility; layout coalescing carries a runtime tail predicate # and therefore explicitly requests ceil-div. coalesce_factor = int(getattr(metadata, "coalesce_factor", 1) or 1) coalesce_axis = int(getattr(metadata, "coalesce_axis", -1)) diff --git a/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageCostModels.h b/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageCostModels.h index aa9db382d6..fad8d1e732 100644 --- a/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageCostModels.h +++ b/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageCostModels.h @@ -40,6 +40,7 @@ enum class StageCostModelKind { IndependentPipelinedLoop, LoopCarriedRecurrence, RowwiseReduction, + PrefixScan, CubeRoofline, TinyCubeRoofline, ConversionPack, @@ -119,6 +120,7 @@ struct StageModeProfile { double storeWarpInstructionsPerCycle = 0.0; double predicateOperationsPerCycle = 0.0; double shuffleLanesPerCycle = 0.0; + double prefixScanDependencyFactor = 1.0; double dotSetupCycles = 0.0; double dotFlopsPerCycle = 0.0; double scalarOperationsPerCycle = 0.0; diff --git a/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageRouteCostModel.h b/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageRouteCostModel.h index 9de0c43732..9dce4aeb96 100644 --- a/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageRouteCostModel.h +++ b/third_party/ascend/costmodel/include/AscendModel/RouteModel/StageRouteCostModel.h @@ -55,6 +55,7 @@ struct StageModelFeatures { bool hasContiguousMemory = false; bool hasIndirectMemory = false; bool hasReduction = false; + bool hasPrefixScan = false; bool hasDot = false; bool hasConversionPack = false; /// True when the Stage is part of the logical-program body created by diff --git a/third_party/ascend/costmodel/lib/AscendModel/Analysis/SimtAnchorAnalysis.cpp b/third_party/ascend/costmodel/lib/AscendModel/Analysis/SimtAnchorAnalysis.cpp index 13de0f1bbc..666b53f89d 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/Analysis/SimtAnchorAnalysis.cpp +++ b/third_party/ascend/costmodel/lib/AscendModel/Analysis/SimtAnchorAnalysis.cpp @@ -61,10 +61,6 @@ analyzePlainOneDimensionalCumsum(Operation *op) { int64_t axisValue = axis.getInt(); if (axisValue < 0 || axisValue >= sourceType.getRank()) return std::nullopt; - for (auto [index, extent] : llvm::enumerate(sourceType.getShape())) - if (static_cast(index) != axisValue && extent != 1) - return std::nullopt; - int64_t realCombineOps = 0; bool isAdd = false; if (op->getNumRegions() != 1 || op->getRegion(0).empty()) @@ -505,7 +501,6 @@ static std::optional analyzeAnchor(Operation *op, if (!facts) return std::nullopt; descriptor.kind = SimtAnchorKind::PlainOneDimensionalCumsum; - descriptor.lowerability.allSimd = false; if (facts->axisExtent <= 0 || !isSupportedCumsumType(facts->elementType)) descriptor.lowerability.mixed = false; } else if (name == "tt.atomic_rmw" || name == "tt.atomic_cas") { diff --git a/third_party/ascend/costmodel/lib/AscendModel/Analysis/StagePartitioner.cpp b/third_party/ascend/costmodel/lib/AscendModel/Analysis/StagePartitioner.cpp index c30acd77a4..2909abbff4 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/Analysis/StagePartitioner.cpp +++ b/third_party/ascend/costmodel/lib/AscendModel/Analysis/StagePartitioner.cpp @@ -509,7 +509,9 @@ static StageCostModelKind classifySemanticRoot(Operation *root) { return StageCostModelKind::AutoBlockifyDispatch; if (operationTreeHasTrueLoopCarriedDependency(root)) return StageCostModelKind::LoopCarriedRecurrence; - if (operationTreeHasAnyName(root, {"tt.reduce", "tt.scan", "linalg.reduce"})) + if (operationTreeHasAnyName(root, {"tt.scan"})) + return StageCostModelKind::PrefixScan; + if (operationTreeHasAnyName(root, {"tt.reduce", "linalg.reduce"})) return StageCostModelKind::RowwiseReduction; if (operationTreeHasAnyName(root, {"tt.dot"})) return StageCostModelKind::CubeRoofline; @@ -804,6 +806,7 @@ static int semanticKindPriority(StageCostModelKind kind) { case StageCostModelKind::LoopCarriedRecurrence: return 90; case StageCostModelKind::RowwiseReduction: + case StageCostModelKind::PrefixScan: return 80; case StageCostModelKind::CubeRoofline: case StageCostModelKind::TinyCubeRoofline: @@ -1073,6 +1076,7 @@ llvm::Error StageFeatureAnalysis::analyze(StagePartition &partition) const { } facts.hasReduction |= name == "tt.reduce" || name == "tt.scan" || name == "linalg.reduce"; + facts.hasPrefixScan |= name == "tt.scan"; facts.hasDot |= name == "tt.dot" || name.contains("matmul") || name.contains("mmad"); facts.hasConversionPack |= @@ -1115,6 +1119,8 @@ llvm::Error StageKindClassifier::analyze(StagePartition &partition, return facts.hasLoop && !facts.hasLoopCarriedDataDependency; case StageCostModelKind::RowwiseReduction: return facts.hasReduction; + case StageCostModelKind::PrefixScan: + return facts.hasPrefixScan; case StageCostModelKind::CubeRoofline: case StageCostModelKind::TinyCubeRoofline: return facts.hasDot; @@ -1147,6 +1153,8 @@ llvm::Error StageKindClassifier::analyze(StagePartition &partition, auto derive = [&]() { if (facts.hasLoopCarriedDataDependency) return StageCostModelKind::LoopCarriedRecurrence; + if (facts.hasPrefixScan) + return StageCostModelKind::PrefixScan; if (facts.hasReduction) return StageCostModelKind::RowwiseReduction; if (facts.hasDot) diff --git a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/SimdSimtCostModel.cpp b/third_party/ascend/costmodel/lib/AscendModel/RouteModel/SimdSimtCostModel.cpp index ebe60f6759..2adb293a4b 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/SimdSimtCostModel.cpp +++ b/third_party/ascend/costmodel/lib/AscendModel/RouteModel/SimdSimtCostModel.cpp @@ -243,6 +243,9 @@ static void readStageResources(ProfileJSONReader &reader, reader.number(*resources, "issue_instructions_per_system_cycle", prefix); profile.spillTransactionsPerCycle = reader.number(*resources, "spill_transactions_per_system_cycle", prefix); + if (const auto *scan = resources->getObject("prefix_scan")) + profile.prefixScanDependencyFactor = + reader.number(*scan, "dependency_factor", prefix + ".prefix_scan"); if (const auto *indirect = reader.object(*resources, "indirect_memory", prefix)) { const std::string path = prefix + ".indirect_memory"; @@ -469,11 +472,11 @@ loadCandidateProfile(llvm::StringRef requestedPath) { return llvm::createStringError( std::errc::invalid_argument, "invalid SIMD/SIMT profile '%s': %s", path.c_str(), reader.getError().str().c_str()); - if (hardware.profileVersion != "david-v100-simd-simt-20260824-v19") + if (hardware.profileVersion != "david-v100-simd-simt-20260903-v20") return llvm::createStringError( std::errc::invalid_argument, "unsupported SIMD/SIMT profile version '%s' " - "(expected david-v100-simd-simt-20260824-v19)", + "(expected david-v100-simd-simt-20260903-v20)", hardware.profileVersion.c_str()); if (!microbench) return llvm::createStringError(std::errc::invalid_argument, diff --git a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageCostModels.cpp b/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageCostModels.cpp index d693c892a1..6a687f44d9 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageCostModels.cpp +++ b/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageCostModels.cpp @@ -262,6 +262,17 @@ static double estimateStage(const LogicalStage &stage, count * std::max(r.scalar + r.load + r.store + r.criticalPath + controlBody(r) + r.spill, r.issue); + case StageCostModelKind::PrefixScan: { + const double scanCritical = + r.compute + r.predicate + + r.shuffle * (mode == StageMode::SIMD + ? profile.simd.prefixScanDependencyFactor + : profile.simt.prefixScanDependencyFactor); + return r.setup + + count * std::max(r.scalar + r.load + r.store + scanCritical + + controlBody(r) + r.spill, + r.issue); + } case StageCostModelKind::CubeRoofline: case StageCostModelKind::TinyCubeRoofline: if (mode == StageMode::SIMD && permitsSimdOverlap(stage)) @@ -335,6 +346,8 @@ llvm::StringRef mlir::ascend::stringifyStageCostModel(StageCostModelKind kind) { return "loop_carried_recurrence"; case StageCostModelKind::RowwiseReduction: return "rowwise_reduction"; + case StageCostModelKind::PrefixScan: + return "prefix_scan"; case StageCostModelKind::CubeRoofline: return "cube_roofline"; case StageCostModelKind::TinyCubeRoofline: @@ -355,7 +368,7 @@ bool StageControlFlowRates::isFiniteAndNonNegative() const { } bool StageModeProfile::isValid(StageMode mode) const { - const std::array common = {setupCycles, + const std::array common = {setupCycles, predicateOperationsPerCycle, shuffleLanesPerCycle, dotSetupCycles, @@ -365,6 +378,7 @@ bool StageModeProfile::isValid(StageMode mode) const { spillTransactionsPerCycle, indirectLoadTransactionsPerCycle, indirectStoreTransactionsPerCycle, + prefixScanDependencyFactor, static_cast(vectorWidth), static_cast(issueWidth)}; if (!std::all_of( diff --git a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageRouteCostModel.cpp b/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageRouteCostModel.cpp index 539a683252..d65d107336 100644 --- a/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageRouteCostModel.cpp +++ b/third_party/ascend/costmodel/lib/AscendModel/RouteModel/StageRouteCostModel.cpp @@ -205,6 +205,7 @@ llvm::json::Object StageModelFeatures::toJSON() const { result["has_contiguous_memory"] = hasContiguousMemory; result["has_indirect_memory"] = hasIndirectMemory; result["has_reduction"] = hasReduction; + result["has_prefix_scan"] = hasPrefixScan; result["has_dot"] = hasDot; result["has_conversion_pack"] = hasConversionPack; result["replicated_by_local_superblock"] = replicatedByLocalSuperBlock; diff --git a/third_party/ascend/costmodel/profiles/simd_simt/david_v100_simd_simt_v1.json b/third_party/ascend/costmodel/profiles/simd_simt/david_v100_simd_simt_v1.json index ca6e30556a..d029927a2c 100644 --- a/third_party/ascend/costmodel/profiles/simd_simt/david_v100_simd_simt_v1.json +++ b/third_party/ascend/costmodel/profiles/simd_simt/david_v100_simd_simt_v1.json @@ -1,7 +1,7 @@ { "$schema": "./simd_simt_profile_schema.json", "schema_version": 10, - "profile_version": "david-v100-simd-simt-20260824-v19", + "profile_version": "david-v100-simd-simt-20260903-v20", "target": "Ascend950PR/dav-c310", "description": "Stage Route Model hardware profile. Stage partitions and workloads are generated online from transformed TTIR; this file contains no workload-specific Stage template or route multiplier.", "microbenchmark_profile": "../microbench/ascend_davidv100_v1.json", @@ -138,6 +138,10 @@ "scalar_operations_per_system_cycle": 1.0, "issue_instructions_per_system_cycle": 6.0, "spill_transactions_per_system_cycle": 1.0, + "prefix_scan": { + "dependency_factor": 17.35, + "description": "SIMD prefix-scan dependency factor calibrated from the identity-baseline-subtracted A5 batch-8, extent-16 scan sweep (143.398 us SIMD scan increment versus 8.267 us SIMT, ratio 17.345); unlike reduction throughput, successive scan levels remain dependent." + }, "indirect_memory": { "load_transactions_per_system_cycle": 0.125, "store_transactions_per_system_cycle": 0.125, @@ -272,6 +276,10 @@ "scalar_operations_per_system_cycle": 4.0, "issue_instructions_per_system_cycle": 4.0, "spill_transactions_per_system_cycle": 1.0, + "prefix_scan": { + "dependency_factor": 1.0, + "description": "SIMT shuffle throughput already represents the lane dependency chain measured by the A5 scan sweep." + }, "indirect_memory": { "load_transactions_per_system_cycle": 0.5, "store_transactions_per_system_cycle": 0.5, diff --git a/third_party/ascend/costmodel/profiles/simd_simt/simd_simt_profile_schema.json b/third_party/ascend/costmodel/profiles/simd_simt/simd_simt_profile_schema.json index bb6d739e23..a6e5487190 100644 --- a/third_party/ascend/costmodel/profiles/simd_simt/simd_simt_profile_schema.json +++ b/third_party/ascend/costmodel/profiles/simd_simt/simd_simt_profile_schema.json @@ -467,6 +467,24 @@ "type": "number", "exclusiveMinimum": 0 }, + "prefix_scan": { + "type": "object", + "required": [ + "dependency_factor", + "description" + ], + "properties": { + "dependency_factor": { + "type": "number", + "exclusiveMinimum": 0 + }, + "description": { + "type": "string", + "minLength": 1 + } + }, + "additionalProperties": false + }, "indirect_memory": { "type": "object", "required": [ diff --git a/third_party/ascend/lib/TritonToLinalg/TileChunkCoalescing.cpp b/third_party/ascend/lib/TritonToLinalg/TileChunkCoalescing.cpp index f7a2f2b82a..1593f0699a 100644 --- a/third_party/ascend/lib/TritonToLinalg/TileChunkCoalescing.cpp +++ b/third_party/ascend/lib/TritonToLinalg/TileChunkCoalescing.cpp @@ -89,7 +89,31 @@ static bool isLiftable(Operation *op) { } return isa(op); + triton::ScanOp, triton::ReduceOp, triton::AssertOp>(op); +} + +// Integer overflow checks emitted by the Triton frontend form predicate-only +// chains such as cmp + cmp -> and -> tt.assert. They constrain no memory +// access and remain valid when lifted to the added coalescing dimension. +static bool isAssertOnlyPredicate(Value value, + DenseSet &visitedPredicates) { + if (!visitedPredicates.insert(value).second) + return true; + if (value.use_empty()) + return false; + for (Operation *user : value.getUsers()) { + if (isa(user)) + continue; + auto andOp = dyn_cast(user); + if (!andOp || !isAssertOnlyPredicate(andOp.getResult(), visitedPredicates)) + return false; + } + return true; +} + +static bool isAssertOnlyPredicate(Value value) { + DenseSet visitedPredicates; + return isAssertOnlyPredicate(value, visitedPredicates); } // Detect tile-index signature: @@ -175,6 +199,8 @@ static std::optional findSeed(ModuleOp moduleOp) { mask = cmp.getResult(); continue; } + if (isAssertOnlyPredicate(cmp.getResult())) + continue; unsafe = true; break; } diff --git a/third_party/ascend/unittest/costmodel_ut/SimdSimtCostModelTest.cpp b/third_party/ascend/unittest/costmodel_ut/SimdSimtCostModelTest.cpp index a5e6d33c9f..74a8fa2154 100644 --- a/third_party/ascend/unittest/costmodel_ut/SimdSimtCostModelTest.cpp +++ b/third_party/ascend/unittest/costmodel_ut/SimdSimtCostModelTest.cpp @@ -344,6 +344,29 @@ TEST(SimdSimtCostModelTest, MixedScopeSuperBlockAmortizesOnlyFixedTransitions) { EXPECT_DOUBLE_EQ(routes->mixed.entryTransitionCycles[1], 108.0); } +TEST(SimdSimtCostModelTest, PrefixScanUsesModeSpecificDependencyFactor) { + LogicalStage stage = logicalStage("scan", StageCostModelKind::PrefixScan); + stage.features.hasReduction = true; + stage.features.hasPrefixScan = true; + stage.workload.operationElements.clear(); + stage.workload.scalarOperations = 0.0; + stage.workload.issueElements = 64.0; + stage.workload.shuffleLaneSteps = 320.0; + + HardwareProfile profile = hardwareProfile(); + profile.simd.prefixScanDependencyFactor = 2.5; + profile.simt.prefixScanDependencyFactor = 1.0; + auto table = evaluateOneStage(std::move(stage), profile); + if (!table) + FAIL() << llvm::toString(table.takeError()); + + const auto &implementations = table->stages.front().implementations; + ASSERT_EQ(implementations.size(), 2u); + EXPECT_EQ(implementations[0].implementation.mode, StageMode::SIMD); + EXPECT_EQ(implementations[1].implementation.mode, StageMode::SIMT); + EXPECT_GT(implementations[0].totalCycles, implementations[1].totalCycles); +} + TEST(SimdSimtCostModelTest, IndependentLoopUsesSimdRooflineAndSerialSimtCost) { LogicalStage stage = logicalStage("independent", StageCostModelKind::IndependentPipelinedLoop, diff --git a/third_party/ascend/unittest/pytest_ut/test_auto_blockify_v1_policy.py b/third_party/ascend/unittest/pytest_ut/test_auto_blockify_v1_policy.py index 7d1eeced77..047d9dae6e 100644 --- a/third_party/ascend/unittest/pytest_ut/test_auto_blockify_v1_policy.py +++ b/third_party/ascend/unittest/pytest_ut/test_auto_blockify_v1_policy.py @@ -91,8 +91,8 @@ def test_route_transform_capability_is_single_resolved_fact(): } opt = SimpleNamespace(compile_on_910_95=True, num_warps=4, logical_program_count_hint=9) capability = __import__("json").loads(_publish_route_transform_capability(metadata, opt)) - assert capability["row_coalescing_applied"] - assert capability["row_coalescing_factor"] == 8 + assert capability["layout_coalescing_applied"] + assert capability["layout_coalescing_factor"] == 8 assert capability["auto_blockify_v1_materializable"] assert capability["whole_kernel_superblock_factors"] == [1, 2, 4] assert capability["scope_superblock_factors"] == [1, 2, 4] diff --git a/third_party/ascend/unittest/pytest_ut/test_launcher_export_api.py b/third_party/ascend/unittest/pytest_ut/test_launcher_export_api.py index be5936c6ae..4860e2d2d3 100644 --- a/third_party/ascend/unittest/pytest_ut/test_launcher_export_api.py +++ b/third_party/ascend/unittest/pytest_ut/test_launcher_export_api.py @@ -104,7 +104,7 @@ def test_make_launcher_shrinks_coalesced_grid_for_both_launch_paths( @patch.object(driver, "is_ffts_supported", return_value=True) @patch.object(driver, "get_ascend_arch_from_env", return_value="Ascend910B") @patch.object(driver, "get_backend_func", side_effect=_mock_backend_func) -def test_make_launcher_ceil_divides_row_coalesced_grid( +def test_make_launcher_ceil_divides_layout_coalesced_grid( _mock_backend_func_patch, _mock_arch, _mock_ffts, diff --git a/third_party/ascend/unittest/pytest_ut/test_simd_simt_costmodel_cases.py b/third_party/ascend/unittest/pytest_ut/test_simd_simt_costmodel_cases.py index 395bd4e98f..f055772f0b 100644 --- a/third_party/ascend/unittest/pytest_ut/test_simd_simt_costmodel_cases.py +++ b/third_party/ascend/unittest/pytest_ut/test_simd_simt_costmodel_cases.py @@ -11,6 +11,11 @@ import triton.runtime.driver as driver from triton.backends.ascend.utils import is_compile_on_910_95 +pytestmark = pytest.mark.skipif( + os.getenv("TRITON_RUN_SIMD_SIMT_COSTMODEL_GUARDS") != "1", + reason="SIMD/SIMT costmodel performance guards are opt-in", +) + simd_simt_910_95_only = pytest.mark.xfail( not is_compile_on_910_95(), reason="SIMD/SIMT cost model only supports 910_95", @@ -474,10 +479,10 @@ def launch(): capability = report["route_transform_capability"] assert capability["source_logical_program_count_hint"] == valid if layout_merge_disabled: - assert not capability["row_coalescing_applied"] + assert not capability["layout_coalescing_applied"] assert capability["logical_program_count_hint"] == valid else: - assert capability["row_coalescing_factor"] == 2 + assert capability["layout_coalescing_factor"] == 2 assert capability["logical_program_count_hint"] == valid // 2 documented_us = 20.578 if layout_merge_disabled else 8.904 _assert_performance(