From 53744ca94a44110ee6b7a329b94cc75f2a8ba005 Mon Sep 17 00:00:00 2001 From: zhywang Date: Tue, 28 Jul 2026 23:32:09 -0700 Subject: [PATCH] feat(ptodsl): Add tile debug print support Add PTODSL support for the native `pto.tprint` debug-print operation and expose it through `pto.tile.print`. The operation is modeled as a device-side side effect, emits no numeric result, and is routed through the native EmitC TPRINT lowering path. ## Summary - Add `_ops.tprint(src, *, tmp=None)` and export it through `pto.tile.print` for tile debug printing, including the optional scratch `tmp` operand - Reject `pto.tile.print` on the VPTO backend with an explicit unsupported diagnostic, while allowing EmitC modules to keep the native `pto.tprint` operation - Treat `TPrintOp` as a native operation in tile expansion and template metadata insertion, so it is not expanded through TileLib templates - Build single-backend EmitC PTODSL kernels through `ptoas` C++ emission followed by `bisheng` fatobj compilation before native launch linking - Add JIT compile coverage for EmitC `pto.tile.print` lowering and the VPTO unsupported-backend diagnostic - Add an A5 TileLib ST case for `pto.tile.print` using an EmitC kernel and a tile fill -> print sequence - Add a VPTO lit negative case that checks the explicit unsupported diagnostic - Document `pto.tile.print` in the PTODSL user guide, including side-effect behavior, `tmp` handling, dtype coverage, backend support, and A5 constraints ## Testing - [x] `ptodsl/tests/test_jit_compile.py` adds EmitC lowering and VPTO unsupported-backend checks for `pto.tile.print`, plus native-build coverage for single-backend EmitC fatobj generation - [x] `test/lit/vpto/tprint_vpto_unsupported.pto` covers the VPTO unsupported diagnostic for `pto.tprint` - [x] `test/tilelib-st/a5/tprint/case.py` covers the A5 EmitC `pto.tile.print` path with an f32 tile fill and device-side print --- lib/PTO/Transforms/ExpandTileOp.cpp | 10 ++ .../Transforms/InsertTemplateAttributes.cpp | 2 +- .../user_guide/04-type-system-and-buffer.md | 31 ++++++ ptodsl/ptodsl/_ops.py | 36 ++++++- ptodsl/ptodsl/_runtime/native_build.py | 100 ++++++++++++++++-- ptodsl/ptodsl/_tile_namespace.py | 1 + ptodsl/tests/test_jit_compile.py | 75 ++++++++++++- ptodsl/tests/test_tilelib_catalog.py | 1 - test/lit/vpto/tprint_vpto_unsupported.pto | 15 +++ test/tilelib-st/a5/tprint/case.py | 50 +++++++++ 10 files changed, 309 insertions(+), 12 deletions(-) create mode 100644 test/lit/vpto/tprint_vpto_unsupported.pto create mode 100644 test/tilelib-st/a5/tprint/case.py diff --git a/lib/PTO/Transforms/ExpandTileOp.cpp b/lib/PTO/Transforms/ExpandTileOp.cpp index 5410383e34..61f51ba230 100644 --- a/lib/PTO/Transforms/ExpandTileOp.cpp +++ b/lib/PTO/Transforms/ExpandTileOp.cpp @@ -1348,8 +1348,18 @@ LogicalResult ExpandState::expandTileOpsInFunction(func::FuncOp func, OpBuilder builder(ctx); // Collect tile ops first (avoid modifying while iterating). + WalkResult unsupportedNativeResult = func.walk([&](pto::TPrintOp op) { + op.emitError("ExpandTileOp: pto.tprint is only supported by the " + "EmitC backend; VPTO lowering for TPRINT is not implemented"); + return WalkResult::interrupt(); + }); + if (unsupportedNativeResult.wasInterrupted()) + return failure(); + SmallVector tileOps; func.walk([&](Operation *op) { + // TReshapeOp is a zero-copy alias and is lowered natively by downstream + // backends, so leave it in place rather than expanding it through TileLib. if (isa(op)) return; if (isa(op)) diff --git a/lib/PTO/Transforms/InsertTemplateAttributes.cpp b/lib/PTO/Transforms/InsertTemplateAttributes.cpp index 7b0f9b28e1..20b06c4525 100644 --- a/lib/PTO/Transforms/InsertTemplateAttributes.cpp +++ b/lib/PTO/Transforms/InsertTemplateAttributes.cpp @@ -955,7 +955,7 @@ struct InsertTemplateAttributesPass SmallVector tileOperations; module.walk([&](Operation *operation) { - if (isa(operation)) + if (isa(operation)) return; if (isa(operation)) tileOperations.push_back(operation); diff --git a/ptodsl/docs/user_guide/04-type-system-and-buffer.md b/ptodsl/docs/user_guide/04-type-system-and-buffer.md index 59d82f3ceb..941c0ea8f2 100644 --- a/ptodsl/docs/user_guide/04-type-system-and-buffer.md +++ b/ptodsl/docs/user_guide/04-type-system-and-buffer.md @@ -382,3 +382,34 @@ scratch = pto.alloc_buffer((32,), pto.f32) | `dtype` | Element type of the returned buffer, such as `pto.f32` or `pto.i32`. | The returned value wraps the buffer address together with its allocation metadata: shape, dtype, element type, element count, and byte size. + +## 4.11 Tile Debug Print + +`pto.tile.print` emits a device-side debug print of a tile's contents (the A5 CCE `TPRINT_IMPL` Vec overload — `cce::printf` over the tile). It is a pure side effect: it produces no value and writes nothing back to a tile, so the call result is `None` and the op is not fusible. Use it only for debugging. It is currently supported only by the EmitC backend, where it lowers directly to a `TPRINT` device call and is skipped by tile-op expansion. + +#### `pto.tile.print(src: Tile, *, tmp: View | None = None) -> None` + +**Description**: Prints `src` from device code. A Unified-Buffer (`vec`) tile prints directly; an accumulator tile may pass a scratch GlobalTensor `tmp` to stage the copy to GM before printing. The optional `printFormat` template argument of the CCE overload is not exposed on this surface (the installed MLIR build has no `PTO_PrintFormatAttr` builder), so the default format is used. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `src` | `Tile` | Tile to print. Printable element types are `f32`, `f16`, `i32`, `i16`, `i8` | +| `tmp` | `View` or `None` | Optional scratch GlobalTensor view (Acc-tile path). Default `None` — direct Vec print | + +**Constraints**: + +- **Vec tiles print without `tmp`**: a tile printed without `tmp` must live in the `vec` (UB) address space. +- **`tmp` is for Mat/Acc staging**: `printFormat` is allowed only when `tmp` is present; on A5, `tmp`-based printing is supported for `vec`/`acc` tiles (Mat-tile printing with `tmp` is A2/A3 only). +- **Backend support**: supported by the EmitC backend. VPTO lowering for `TPRINT` is not implemented. +- **Hardware mapping**: executes on the **Vector pipeline** (`PIPE_V`). + +**Example** — load a tile from GM and print it for debugging: + +```python +src_view = pto.make_tensor_view(src_ptr, shape=[rows, cols], strides=[cols, 1]) +src_tile = pto.alloc_tile(shape=[rows, cols], dtype=pto.f32) +pto.tile.load(src_view, src_tile) +pto.tile.print(src_tile) +``` diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index c3349f5459..6b286d38ee 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -192,6 +192,20 @@ def _current_target_arch(): return getattr(current_module_spec, "target_arch", None) +def _current_backend(): + try: + from ._tracing.active import current_session + session = current_session() + except Exception: + return None + if session is None: + return None + current_module_spec = getattr( + session, "current_function_module_spec", session.module_spec + ) + return getattr(current_module_spec, "backend", None) + + def _require_target_arch(surface: str, allowed: set[str]): target = _current_target_arch() if target is None: @@ -3424,6 +3438,26 @@ def _tile_numel(shape, *, context: str): return numel +def tprint(src, *, tmp=None): + """``pto.tprint ins(src, tmp?)`` -- device-side debug print of a tile. + + Pure side effect (``cce::printf``, no numeric result): a Vec tile prints + directly, while an Acc tile may pass a scratch GlobalTensor ``tmp`` to stage + the copy to GM. Only the default print format is exposed, matching the + no-format C++ overload. + """ + backend = _current_backend() + if backend == "vpto": + raise ValueError( + "pto.tile.print is only supported by the EmitC backend; " + "VPTO lowering for TPRINT is not implemented" + ) + _pto.tprint( + unwrap_surface_value(src), + tmp=None if tmp is None else unwrap_surface_value(tmp), + ) + + def treshape(src, *, shape, dtype=None, blayout=None): """``pto.treshape ins(src) -> result``.""" src_value = unwrap_surface_value(src) @@ -6272,7 +6306,7 @@ def import_reserved_buffer(name, *, peer_func): "trowsum", "trowmax", "trowmin", "trowprod", "trowargmax", "trowargmin", "tcolsum", "tcolmax", "tcolmin", "tcolprod", "tcolargmax", "tcolargmin", "tcmp", "tcmps", - "texpands", "treshape", "trowexpand", "tcolexpand", + "texpands", "tprint", "treshape", "trowexpand", "tcolexpand", "trowexpandadd", "trowexpandsub", "trowexpandmul", "trowexpanddiv", "trowexpandmax", "trowexpandmin", "trowexpandexpdif", "tcolexpandadd", "tcolexpandsub", "tcolexpandmul", "tcolexpanddiv", "tcolexpandmax", "tcolexpandmin", "tcolexpandexpdif", "tsort32", "tmrgsort", "tgather", diff --git a/ptodsl/ptodsl/_runtime/native_build.py b/ptodsl/ptodsl/_runtime/native_build.py index 310b44067e..5ba2b87330 100644 --- a/ptodsl/ptodsl/_runtime/native_build.py +++ b/ptodsl/ptodsl/_runtime/native_build.py @@ -86,12 +86,17 @@ def _source_ptoas_overrides(module_spec) -> dict: return {"backend": module_spec.backend} +def _is_single_emitc_native_build(*, module_spec, mlir_text: str) -> bool: + return module_spec.backend == "emitc" and 'pto.backend = "vpto"' not in mlir_text + + def _compile_config_text( *, module_spec, effective_insert_sync: bool, effective_pto_level: str | None, ptoas_overrides: dict, + emitc_fatobj: bool, ) -> str: return "\n".join( [ @@ -101,6 +106,7 @@ def _compile_config_text( f"insert_sync={effective_insert_sync}", f"pto_level={effective_pto_level}", f"backend={ptoas_overrides.get('backend')}", + f"emitc_fatobj={emitc_fatobj}", "enable_tile_op_expand=True", ] ) @@ -148,6 +154,62 @@ def _kernel_compile_flags(kernel_kind: str, target_arch: str) -> list[str]: ] +def _emitc_fatobj_arch(*, kernel_kind: str, target_arch: str) -> str: + arch = aicore_arch_for_kernel_kind(kernel_kind, target_arch) + for suffix in ("-vec", "-cube"): + if arch.endswith(suffix): + return arch[: -len(suffix)] + return arch + + +def _emitc_fatobj_compile_flags(kernel_kind: str, target_arch: str) -> list[str]: + arch = _emitc_fatobj_arch(kernel_kind=kernel_kind, target_arch=target_arch) + return [ + "-xcce", + "-fenable-matrix", + "--cce-aicore-enable-tl", + "-fPIC", + "-Xhost-start", + "-Xhost-end", + "-mllvm", + "-cce-aicore-stack-size=0x8000", + "-mllvm", + "-cce-aicore-function-stack-size=0x8000", + "-mllvm", + "-cce-aicore-record-overflow=true", + "-mllvm", + "-cce-aicore-addr-transform", + "-mllvm", + "-cce-aicore-dcci-insert-for-scalar=false", + f"--cce-aicore-arch={arch}", + "-DREGISTER_BASE", + "-std=c++17", + "-O2", + "-dc", + *common_include_flags(), + ] + + +def _compile_emitc_cpp_to_fatobj( + kernel_cpp: Path, + kernel_object: Path, + *, + kernel_kind: str, + target_arch: str, +) -> None: + bisheng = resolve_bisheng() + _run( + [ + bisheng, + *_emitc_fatobj_compile_flags(kernel_kind, target_arch), + "-c", + str(kernel_cpp), + "-o", + str(kernel_object), + ] + ) + + def _compile_launch_cpp( launch_cpp: Path, launch_object: Path, @@ -219,11 +281,16 @@ def build_native_library( ) effective_pto_level = _effective_pto_level(mode=module_spec.mode) ptoas_overrides = _source_ptoas_overrides(module_spec) + emitc_fatobj = _is_single_emitc_native_build( + module_spec=module_spec, + mlir_text=mlir_text, + ) compile_config_text = _compile_config_text( module_spec=module_spec, effective_insert_sync=effective_insert_sync, effective_pto_level=effective_pto_level, ptoas_overrides=ptoas_overrides, + emitc_fatobj=emitc_fatobj, ) sim_mode = bool(os.environ.get("MSPROF_SIMULATOR_MODE")) link_config_text = "\n".join(runtime_library_flags(sim_mode=sim_mode)) @@ -241,14 +308,31 @@ def build_native_library( artifacts.mlir_path.write_text(mlir_text, encoding="utf-8") artifacts.launch_cpp.write_text(launch_cpp_text, encoding="utf-8") - _run_ptoas( - artifacts.mlir_path, - artifacts.kernel_object, - target_arch=module_spec.target_arch, - insert_sync=effective_insert_sync, - pto_level=effective_pto_level, - **ptoas_overrides, - ) + if emitc_fatobj: + kernel_cpp = artifacts.cache_dir / "kernel.cpp" + _run_ptoas( + artifacts.mlir_path, + kernel_cpp, + target_arch=module_spec.target_arch, + insert_sync=effective_insert_sync, + pto_level=effective_pto_level, + **ptoas_overrides, + ) + _compile_emitc_cpp_to_fatobj( + kernel_cpp, + artifacts.kernel_object, + kernel_kind=module_spec.kernel_kind, + target_arch=module_spec.target_arch, + ) + else: + _run_ptoas( + artifacts.mlir_path, + artifacts.kernel_object, + target_arch=module_spec.target_arch, + insert_sync=effective_insert_sync, + pto_level=effective_pto_level, + **ptoas_overrides, + ) launch_object = artifacts.cache_dir / "launch.o" export_macro = f"{ir_function_name}_EXPORTS" diff --git a/ptodsl/ptodsl/_tile_namespace.py b/ptodsl/ptodsl/_tile_namespace.py index 0ceb3c534b..a374256ca2 100644 --- a/ptodsl/ptodsl/_tile_namespace.py +++ b/ptodsl/ptodsl/_tile_namespace.py @@ -155,6 +155,7 @@ def rowargmin(src, dst, *, tmp=None): cmps = staticmethod(_ops.tcmps) expands = staticmethod(_ops.texpands) + print = staticmethod(_ops.tprint) reshape = staticmethod(_ops.treshape) rowexpand = staticmethod(_ops.trowexpand) colexpand = staticmethod(_ops.tcolexpand) diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 7047c31b11..386277cc98 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -654,6 +654,18 @@ def tile_surface_compute_probe(): _ = reshape_col +@pto.jit(target="a5") +def tile_print_vpto_backend_probe(): + src = pto.alloc_tile(shape=[1, 16], dtype=pto.f32) + pto.tile.print(src) + + +@pto.jit(target="a5", backend="emitc") +def tile_print_emitc_backend_probe(): + src = pto.alloc_tile(shape=[1, 16], dtype=pto.f32) + pto.tile.print(src) + + @pto.jit(target="a5") def tile_surface_window_matmul_probe(): src_mat = pto.alloc_tile( @@ -4608,12 +4620,14 @@ def inline_source_backed_probe(ptr: pto.ptr(pto.f32, "gm"), rows: pto.i32): ), ("explicit-level3-container", host_vec_copy_explicit_addr.compile(), None), ("same-backend-multi-child-container", kernel_module_compiled, None), + ("single-emitc-container", host_vec_copy_emitc.compile(), None), ("mixed-backend-container", emitc_entry_calls_vpto_kernel_module_probe.compile(), None), ("source-auto", source_native_build_compiled, None), ("source-explicit", source_explicit_native_build_compiled, None), ("source-no-insert-sync", source_no_insert_sync_native_build_compiled, None), ) native_build_observations = [] + emitc_fatobj_observations = [] launch_target_arches = [] with TemporaryDirectory() as tmpdir: @@ -4642,7 +4656,35 @@ def fake_run_ptoas(mlir_path, kernel_object, *, target_arch, insert_sync=None, b "mlir_text": mlir_path.read_text(encoding="utf-8"), } ) - kernel_object.write_text("fake fatobj\n", encoding="utf-8") + if kernel_object.suffix == ".cpp": + kernel_object.write_text("fake emitc c++\n", encoding="utf-8") + else: + kernel_object.write_text("fake fatobj\n", encoding="utf-8") + + def fake_compile_emitc_cpp_to_fatobj( + kernel_cpp, + kernel_object, + *, + kernel_kind, + target_arch, + ): + expect( + kernel_cpp.is_file(), + "single-backend EmitC native build should materialize kernel.cpp", + ) + expect( + kernel_kind in {"vector", "cube"}, + "single-backend EmitC native build should forward kernel kind", + ) + emitc_fatobj_observations.append( + { + "kernel_cpp": kernel_cpp, + "kernel_object": kernel_object, + "kernel_kind": kernel_kind, + "target_arch": target_arch, + } + ) + kernel_object.write_text("fake emitc fatobj\n", encoding="utf-8") def fake_compile_launch_cpp( launch_cpp, @@ -4667,6 +4709,10 @@ def fake_link_shared_library(launch_object, kernel_object, shared_library, *, ke with mock.patch.object(native_build_runtime, "artifact_paths", side_effect=fake_artifacts), mock.patch.object( native_build_runtime, "is_native_build_current", return_value=False ), mock.patch.object(native_build_runtime, "_run_ptoas", side_effect=fake_run_ptoas), mock.patch.object( + native_build_runtime, + "_compile_emitc_cpp_to_fatobj", + side_effect=fake_compile_emitc_cpp_to_fatobj, + ), mock.patch.object( native_build_runtime, "_compile_launch_cpp", side_effect=fake_compile_launch_cpp ), mock.patch.object( native_build_runtime, "_link_shared_library", side_effect=fake_link_shared_library @@ -4730,6 +4776,19 @@ def fake_link_shared_library(launch_object, kernel_object, shared_library, *, ke observation["mlir_text"].count("module") >= 2, f"{label} native build should route the unified outer+child container through ptoas", ) + if label == "single-emitc-container": + expect( + observation["kernel_object"].name == "kernel.cpp", + "single-backend EmitC native build should request C++ text from ptoas before fatobj compilation", + ) + expect( + len(emitc_fatobj_observations) == 1, + "native build should compile exactly the single-backend EmitC ptoas C++ output into a fatobj", + ) + expect( + emitc_fatobj_observations[0]["target_arch"] == "a5", + "single-backend EmitC fatobj compilation should preserve target arch", + ) with TemporaryDirectory() as tmpdir: tmpdir_path = Path(tmpdir) mlir_path = tmpdir_path / "kernel.mlir" @@ -4910,6 +4969,20 @@ def fake_run_ptoas_cmd(cmd, *, cwd=None): expect(tile_sort_gather_text.count("pto.tgather") == 2, "tile gather wrappers should lower to pto.tgather") expect("#pto.mask_pattern" in tile_sort_gather_text, "pto.tile.gather should preserve P0101") expect("#pto.mask_pattern" in tile_sort_gather_text, "pto.tgather should preserve P1010") + expect_raises( + ValueError, + lambda: tile_print_vpto_backend_probe.compile().mlir_text(), + "pto.tile.print is only supported by the EmitC backend", + ) + tile_print_emitc_text = tile_print_emitc_backend_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify( + tile_print_emitc_text, + "tile print EmitC backend specialization", + ) + expect( + "pto.tprint" in tile_print_emitc_text, + "pto.tile.print should lower to pto.tprint on the EmitC backend", + ) tile_ci_text = tile_ci_surface_probe.compile().mlir_text() expect_parse_roundtrip_and_verify(tile_ci_text, "tile ci surface specialization") expect(tile_ci_text.count("pto.tci") == 2, "pto.tile.ci should lower to pto.tci") diff --git a/ptodsl/tests/test_tilelib_catalog.py b/ptodsl/tests/test_tilelib_catalog.py index 2217849fd9..370fd98c8f 100644 --- a/ptodsl/tests/test_tilelib_catalog.py +++ b/ptodsl/tests/test_tilelib_catalog.py @@ -1514,6 +1514,5 @@ def test_tdequant_dtype_versions_render(self): self.assertEqual(selected.name, expected_name) self.assertIn("pto.vmul", selected.specialize(**specs).mlir_text()) - if __name__ == "__main__": unittest.main() diff --git a/test/lit/vpto/tprint_vpto_unsupported.pto b/test/lit/vpto/tprint_vpto_unsupported.pto new file mode 100644 index 0000000000..1d21b97ace --- /dev/null +++ b/test/lit/vpto/tprint_vpto_unsupported.pto @@ -0,0 +1,15 @@ +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto %s -o - 2>&1 | FileCheck %s + +module attributes {pto.backend = "vpto", pto.target_arch = "a5"} { + func.func @tprint_vpto_unsupported() attributes {pto.entry} { + %tile = pto.alloc_tile + : !pto.tile_buf + pto.tprint ins(%tile : !pto.tile_buf) + return + } +} + +// CHECK: ExpandTileOp: pto.tprint is only supported by the EmitC backend; VPTO lowering for TPRINT is not implemented diff --git a/test/tilelib-st/a5/tprint/case.py b/test/tilelib-st/a5/tprint/case.py new file mode 100644 index 0000000000..8c7b740d3b --- /dev/null +++ b/test/tilelib-st/a5/tprint/case.py @@ -0,0 +1,50 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +# PTODSL ST for pto.tprint. TPRINT is a device-side debug side effect, so the +# case validates compile/launch without a numeric output comparison. + +from pathlib import Path +import sys + +if __package__ in {None, ""}: + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from common import auto_main +from ptodsl import pto + + +@pto.jit(name="tprint_f32_16x16", target="a5", backend="emitc") +def _tprint_f32_16x16_kernel(): + tile = pto.alloc_tile(shape=[16, 16], dtype=pto.f32) + + pto.tile.expands(1.0, tile) + pto.tile.print(tile) + + +def _make_case(): + return [], None + + +def _check_case(device_inputs, golden): + _ = device_inputs + _ = golden + + +CASES = [ + { + "name": "tprint_f32_16x16", + "kernel": _tprint_f32_16x16_kernel, + "make_case": _make_case, + "check": _check_case, + } +] + + +auto_main(globals())