From 1df1a2243f53f79819068cbacaf222b54dcd1169 Mon Sep 17 00:00:00 2001 From: zhywang Date: Thu, 6 Aug 2026 03:59:32 -0700 Subject: [PATCH] feat(codegen): Add EmitC tprint fatobj support ## Summary - Add PTO IR, C API, Python binding, and ptodsl frontend support for `tprint` print formats and temporary buffers - Lower `tprint` through EmitC with the required TPrint headers, format mapping, temporary tensor view handling, and S-pipe synchronization - Route ptodsl native builds through `ptoas --fatobj` and select the correct CCE target CPU for EmitC and mixed backend child jobs - Reject unsupported VPTO `tprint` lowering and preserve VPTO ABI imports for mixed backend child modules - Add lit coverage, ptodsl compile coverage, TileLib ST tprint cases, and user-guide documentation for supported tprint behavior ## Testing - [x] PTOASPythonPackage builds locally - [x] `ptodsl/tests/test_jit_compile.py` passes locally - [x] `test/tilelib-st/a5/tprint/case.py --emit-mlir` compiles to fatobj locally with `ptoas --fatobj --enable-tile-op-expand` - [x] `git diff --check` passes --- include/PTO/IR/PTOOps.td | 5 +- include/PTO/Transforms/TileOpExpansionUtils.h | 2 +- include/pto-c/Dialect/PTO.h | 3 + lib/Bindings/Python/PTOModule.cpp | 9 + lib/CAPI/Dialect/PTO.cpp | 1 + lib/PTO/IR/PTO.cpp | 2 + lib/PTO/Transforms/ExpandTileOp.cpp | 15 + lib/PTO/Transforms/PTOToEmitC.cpp | 102 +++++-- .../user_guide/04-type-system-and-buffer.md | 32 ++ ptodsl/ptodsl/_ops.py | 64 +++- ptodsl/ptodsl/_runtime/native_build.py | 12 + ptodsl/ptodsl/_surface_types.py | 2 + ptodsl/ptodsl/_tile_namespace.py | 1 + ptodsl/ptodsl/pto.py | 1 + ptodsl/tests/test_jit_compile.py | 87 +++++- python/pto/dialects/pto.py | 4 + ...backend_single_emitc_child_tprint_view.pto | 29 ++ test/lit/pto/emitc_fatobj_driver_invalid.pto | 8 + test/lit/pto/tprint_alloc_tile_no_rebind.pto | 12 +- ...to => tprint_format_without_tmp_emitc.pto} | 6 +- .../pto/tprint_insert_sync_pipe_selection.pto | 32 ++ test/lit/vpto/tprint_vpto_unsupported.pto | 15 + test/tilelib-st/a5/tprint/case.py | 172 +++++++++++ tools/ptoas/ObjectEmission.cpp | 52 ++-- tools/ptoas/ObjectEmission.h | 12 +- tools/ptoas/driver.cpp | 281 ++++++++++++++++-- tools/ptoas/ptoas.h | 2 + 27 files changed, 882 insertions(+), 81 deletions(-) create mode 100644 test/lit/pto/backend_single_emitc_child_tprint_view.pto create mode 100644 test/lit/pto/emitc_fatobj_driver_invalid.pto rename test/lit/pto/{tprint_format_requires_tmp.pto => tprint_format_without_tmp_emitc.pto} (83%) create mode 100644 test/lit/pto/tprint_insert_sync_pipe_selection.pto create mode 100644 test/lit/vpto/tprint_vpto_unsupported.pto create mode 100644 test/tilelib-st/a5/tprint/case.py diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index 69f1e6c149..9e1bd0f4d9 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -6923,7 +6923,8 @@ def TPrintOp: PTO_TOp<"tprint", [ ]> { let summary = "TPRINT: Print the contents of a Tile or GlobalTensor for debugging purposes directly from device code."; let description = [{ - pto-isa overloads support `TPRINT(src)` and `TPRINT(src, tmp)`. + pto-isa debug wrappers support `TPRINT(src)` and + `TPRINT(src, tmp)`. The optional tmp operand is used when printing Mat/Acc tiles through a scratch GlobalTensor, while Vec tiles and GlobalTensors can be printed without tmp. @@ -6941,7 +6942,7 @@ def TPrintOp: PTO_TOp<"tprint", [ let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ - ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } + ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_S; } }]; } diff --git a/include/PTO/Transforms/TileOpExpansionUtils.h b/include/PTO/Transforms/TileOpExpansionUtils.h index f38b6ca191..4c7e635044 100644 --- a/include/PTO/Transforms/TileOpExpansionUtils.h +++ b/include/PTO/Transforms/TileOpExpansionUtils.h @@ -19,7 +19,7 @@ namespace mlir::pto { inline bool isTileLibExpandableOp(Operation *op) { if (!op || !isa(op)) return false; - return !isa(op); } diff --git a/include/pto-c/Dialect/PTO.h b/include/pto-c/Dialect/PTO.h index cb653c2384..d996029e99 100644 --- a/include/pto-c/Dialect/PTO.h +++ b/include/pto-c/Dialect/PTO.h @@ -160,6 +160,9 @@ MLIR_CAPI_EXPORTED int32_t mlirPTOSqrtPrecisionAttrGetValue(MlirAttribute attr); MLIR_CAPI_EXPORTED MlirAttribute mlirPTOFmodPrecisionAttrGet(MlirContext ctx, int32_t value); MLIR_CAPI_EXPORTED bool mlirPTOAttrIsAFmodPrecisionAttr(MlirAttribute attr); MLIR_CAPI_EXPORTED int32_t mlirPTOFmodPrecisionAttrGetValue(MlirAttribute attr); +MLIR_CAPI_EXPORTED MlirAttribute mlirPTOPrintFormatAttrGet(MlirContext ctx, int32_t value); +MLIR_CAPI_EXPORTED bool mlirPTOAttrIsAPrintFormatAttr(MlirAttribute attr); +MLIR_CAPI_EXPORTED int32_t mlirPTOPrintFormatAttrGetValue(MlirAttribute attr); MLIR_CAPI_EXPORTED MlirAttribute mlirPTOSaturationModeAttrGet(MlirContext ctx, int32_t value); MLIR_CAPI_EXPORTED bool mlirPTOAttrIsASaturationModeAttr(MlirAttribute attr); MLIR_CAPI_EXPORTED int32_t mlirPTOSaturationModeAttrGetValue(MlirAttribute attr); diff --git a/lib/Bindings/Python/PTOModule.cpp b/lib/Bindings/Python/PTOModule.cpp index 94d3f728f4..4826870638 100644 --- a/lib/Bindings/Python/PTOModule.cpp +++ b/lib/Bindings/Python/PTOModule.cpp @@ -213,6 +213,11 @@ void mlir::pto::python::populatePTODialectBindings(pybind11::module_ &m) { .value("Default", mlir::pto::FmodPrecision::Default) .value("HighPrecision", mlir::pto::FmodPrecision::HighPrecision); + py::enum_(m, "PrintFormat") + .value("Width8_Precision4", mlir::pto::PrintFormat::Width8_Precision4) + .value("Width8_Precision2", mlir::pto::PrintFormat::Width8_Precision2) + .value("Width10_Precision6", mlir::pto::PrintFormat::Width10_Precision6); + py::enum_(m, "SaturationMode") .value("ON", mlir::pto::SaturationMode::ON) .value("OFF", mlir::pto::SaturationMode::OFF); @@ -578,6 +583,10 @@ void mlir::pto::python::populatePTODialectBindings(pybind11::module_ &m) { mlirPTOAttrIsAFmodPrecisionAttr, mlirPTOFmodPrecisionAttrGet, mlirPTOFmodPrecisionAttrGetValue); + bindPTOEnumAttr(m, "PrintFormatAttr", "PrintFormat", + mlirPTOAttrIsAPrintFormatAttr, + mlirPTOPrintFormatAttrGet, + mlirPTOPrintFormatAttrGetValue); mlir_attribute_subclass( m, "SaturationModeAttr", diff --git a/lib/CAPI/Dialect/PTO.cpp b/lib/CAPI/Dialect/PTO.cpp index 40520b90fc..a4f341e561 100644 --- a/lib/CAPI/Dialect/PTO.cpp +++ b/lib/CAPI/Dialect/PTO.cpp @@ -424,6 +424,7 @@ DEFINE_PTO_ENUM_ATTR_CAPI(RemPrecision, RemPrecisionAttr, RemPrecision) DEFINE_PTO_ENUM_ATTR_CAPI(RsqrtPrecision, RsqrtPrecisionAttr, RsqrtPrecision) DEFINE_PTO_ENUM_ATTR_CAPI(SqrtPrecision, SqrtPrecisionAttr, SqrtPrecision) DEFINE_PTO_ENUM_ATTR_CAPI(FmodPrecision, FmodPrecisionAttr, FmodPrecision) +DEFINE_PTO_ENUM_ATTR_CAPI(PrintFormat, PrintFormatAttr, PrintFormat) #undef DEFINE_PTO_ENUM_ATTR_CAPI diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 86907e3255..a5b22cecac 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -12885,10 +12885,12 @@ void mlir::pto::TPrintOp::print(OpAsmPrinter &p) { mlir::LogicalResult mlir::pto::TPrintOp::verify() { auto srcType = getSrc().getType(); Value tmp = getTPrintTmpIfPresent(*this); + auto printFormatAttr = dyn_cast_or_null(getProperties().printFormat); if (printFormatAttr && !tmp) return emitOpError() << "expects printFormat only when tmp is present"; + if (auto tb = mlir::dyn_cast(srcType)) { auto elem = tb.getElementType(); if (!(elem.isF16() || elem.isF32() || diff --git a/lib/PTO/Transforms/ExpandTileOp.cpp b/lib/PTO/Transforms/ExpandTileOp.cpp index 471afcb928..8ce04b36eb 100644 --- a/lib/PTO/Transforms/ExpandTileOp.cpp +++ b/lib/PTO/Transforms/ExpandTileOp.cpp @@ -1260,6 +1260,21 @@ void ExpandTileOpPass::runOnOperation() { ModuleOp mod = getOperation(); MLIRContext *ctx = &getContext(); + bool isVPTOBackend = false; + if (auto backend = mod->getAttrOfType("pto.backend")) + isVPTOBackend = backend.getValue() == "vpto"; + if (isVPTOBackend) { + WalkResult unsupportedTPrint = mod.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 (unsupportedTPrint.wasInterrupted()) { + signalPassFailure(); + return; + } + } + bool hasExpandableOps = false; mod.walk([&](Operation *op) { if (pto::isTileLibExpandableOp(op)) { diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 271f9e50a8..5a8ed7854e 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -13718,34 +13718,82 @@ struct EmitPTOManualPass return signalPassFailure(); } - bool needsEventIdArrayHelper = false; - bool needsTRandomHelper = false; - bool needsGlobalTensorDataHelper = false; - mop.walk([&](Operation *op) { - if (isa(op)) - needsEventIdArrayHelper = true; - if (isa(op)) - needsTRandomHelper = true; - if (auto cmo = dyn_cast(op)) { - if (cmo.getAddr()) - needsGlobalTensorDataHelper = true; - } - if (auto init = dyn_cast(op)) { - if (isa(init.getGmAddr().getType())) - needsGlobalTensorDataHelper = true; - } - if (isa(op)) - needsGlobalTensorDataHelper = true; - }); + bool needsEventIdArrayHelper = false; + bool needsTRandomHelper = false; + bool needsGlobalTensorDataHelper = false; + bool needsTPrintInclude = false; + mop.walk([&](Operation *op) { + if (isa(op)) + needsEventIdArrayHelper = true; + if (isa(op)) + needsTRandomHelper = true; + if (isa(op)) + needsTPrintInclude = true; + if (auto cmo = dyn_cast(op)) { + if (cmo.getAddr()) + needsGlobalTensorDataHelper = true; + } + if (auto init = dyn_cast(op)) { + if (isa(init.getGmAddr().getType())) + needsGlobalTensorDataHelper = true; + } + if (isa(op)) + needsGlobalTensorDataHelper = true; + }); - // 1. 插入头文件 - auto loc = mop->getLoc(); - OpBuilder builder(ctx); - builder.setInsertionPointToStart(mop.getBody()); - builder.create( - loc, "pto/pto-inst.hpp", /*is_standard_include=*/false); - builder.create( - loc, builder.getStringAttr("using namespace pto;")); + auto loc = mop->getLoc(); + OpBuilder builder(ctx); + builder.setInsertionPointToStart(mop.getBody()); + if (needsTPrintInclude) { + // CANN 9.1 asc_printf.h defines a global conditional helper that can + // collide with std::conditional through unqualified lookup. Keep the + // workaround scoped to that include. + builder.create( + loc, "cstdint", /*is_standard_include=*/true); + builder.create( + loc, "type_traits", /*is_standard_include=*/true); + builder.create( + loc, builder.getStringAttr( + "#define conditional PTOAS_ASC_PRINTF_CONDITIONAL")); + builder.create( + loc, "utils/debug/asc_printf.h", + /*is_standard_include=*/false); + builder.create( + loc, builder.getStringAttr("#undef conditional")); + } + builder.create( + loc, "pto/pto-inst.hpp", /*is_standard_include=*/false); + if (needsTPrintInclude) { + builder.create( + loc, builder.getStringAttr(R"cpp(namespace cce { +template +AICORE inline void printf(const __gm__ char *fmt, Args &&...args) { + __asc_aicore::printf(fmt, args...); +} +})cpp")); + builder.create( + loc, + targetArch == PTOArch::A5 ? "pto/npu/a5/TPrint.hpp" + : "pto/npu/a2a3/TPrint.hpp", + /*is_standard_include=*/false); + builder.create( + loc, builder.getStringAttr(R"cpp(#if !defined(_DEBUG) && !defined(__CPU_SIM) +template +AICORE inline void TPRINT(TileData &src) { + TPRINT_IMPL(src); +} + +template +AICORE inline void TPRINT(TileData &src, GlobalData &tmp) { + TPRINT_IMPL(src, tmp); +} +#endif +)cpp")); + } + builder.create( + loc, builder.getStringAttr("using namespace pto;")); // Emit a C++ definition for every !pto.struct used in the module, in // dependency order (nested structs first) so there is no 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 39a9cf9460..fc1d55052b 100644 --- a/ptodsl/docs/user_guide/04-type-system-and-buffer.md +++ b/ptodsl/docs/user_guide/04-type-system-and-buffer.md @@ -412,3 +412,35 @@ 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 through the PTO-ISA `TPRINT` wrapper. 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 native `TPRINT` device call and is skipped by tile-op expansion. + +#### `pto.tile.print(src: Tile, *, tmp: View | None = None, print_format: str | pto.PrintFormat | 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. `print_format` maps to the C++ `TPRINT(...)` template argument and requires `tmp`. Supported string values are `"width8_precision4"` (default), `"width8_precision2"`, and `"width10_precision6"`. + +**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 | +| `print_format` | `str`, `pto.PrintFormat`, or `None` | Optional print format used with `tmp`. Default `None` uses `width8_precision4` | + +**Constraints**: + +- **Vec tiles print without `tmp`**: a tile printed without `tmp` must live in the `vec` (UB) address space. +- **`tmp` is for formatted printing and Mat/Acc staging**: `print_format` requires `tmp`; 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 **Scalar pipeline** (`PIPE_S`). + +**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 013f4936e4..4f6492975f 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: @@ -3625,6 +3639,54 @@ def _tile_numel(shape, *, context: str): return numel +_PRINT_FORMAT_ALIASES = { + "width8_precision4": "Width8_Precision4", + "width8_precision2": "Width8_Precision2", + "width10_precision6": "Width10_Precision6", + "Width8_Precision4": "Width8_Precision4", + "Width8_Precision2": "Width8_Precision2", + "Width10_Precision6": "Width10_Precision6", +} + + +def _coerce_print_format(print_format): + if print_format is None: + return None + if isinstance(print_format, Attribute): + return print_format + if isinstance(print_format, str): + enum_name = _PRINT_FORMAT_ALIASES.get(print_format) + if enum_name is None: + expected = ", ".join(sorted(_PRINT_FORMAT_ALIASES)) + raise ValueError(f"pto.tile.print(print_format=...) expected one of: {expected}") + print_format = getattr(_pto.PrintFormat, enum_name) + return _pto.PrintFormatAttr.get(print_format) + + +def tprint(src, *, tmp=None, print_format=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. ``print_format`` requires ``tmp`` and accepts + ``"width8_precision4"``, ``"width8_precision2"``, + ``"width10_precision6"``, or ``pto.PrintFormat``. + """ + 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" + ) + if print_format is not None and tmp is None: + raise ValueError("pto.tile.print(print_format=...) requires tmp") + _pto.tprint( + unwrap_surface_value(src), + tmp=None if tmp is None else unwrap_surface_value(tmp), + print_format=_coerce_print_format(print_format), + ) + + def treshape(src, *, shape, dtype=None, blayout=None): """``pto.treshape ins(src) -> result``.""" src_value = unwrap_surface_value(src) @@ -6627,7 +6689,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", "tscatter", diff --git a/ptodsl/ptodsl/_runtime/native_build.py b/ptodsl/ptodsl/_runtime/native_build.py index 22516c7eb7..e373364404 100644 --- a/ptodsl/ptodsl/_runtime/native_build.py +++ b/ptodsl/ptodsl/_runtime/native_build.py @@ -47,6 +47,7 @@ def _run_ptoas( insert_sync: bool | None = None, backend: str | None = None, pto_level: str | None = None, + fatobj: bool = False, ) -> None: ptoas = resolve_ptoas_binary() cmd = [ @@ -59,6 +60,8 @@ def _run_ptoas( cmd.append(f"--pto-level={pto_level}") if insert_sync is True: cmd.append("--enable-insert-sync") + if fatobj: + cmd.append("--fatobj") cmd.extend([ "--enable-tile-op-expand", str(mlir_path), @@ -86,12 +89,17 @@ def _source_ptoas_overrides(module_spec) -> dict: return {"backend": module_spec.backend} +def _requires_fatobj(mlir_text: str) -> bool: + return "pto.tprint" in mlir_text + + def _compile_config_text( *, module_spec, effective_insert_sync: bool, effective_pto_level: str | None, ptoas_overrides: dict, + fatobj: bool, ) -> str: return "\n".join( [ @@ -101,6 +109,7 @@ def _compile_config_text( f"insert_sync={effective_insert_sync}", f"pto_level={effective_pto_level}", f"backend={ptoas_overrides.get('backend')}", + f"fatobj={fatobj}", "enable_tile_op_expand=True", ] ) @@ -219,11 +228,13 @@ def build_native_library( ) effective_pto_level = _effective_pto_level(mode=module_spec.mode) ptoas_overrides = _source_ptoas_overrides(module_spec) + fatobj = _requires_fatobj(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, + fatobj=fatobj, ) sim_mode = bool(os.environ.get("MSPROF_SIMULATOR_MODE")) link_config_text = "\n".join(runtime_library_flags(sim_mode=sim_mode)) @@ -247,6 +258,7 @@ def build_native_library( target_arch=module_spec.target_arch, insert_sync=effective_insert_sync, pto_level=effective_pto_level, + fatobj=fatobj, **ptoas_overrides, ) diff --git a/ptodsl/ptodsl/_surface_types.py b/ptodsl/ptodsl/_surface_types.py index 0074c10e62..b1c8398ecb 100644 --- a/ptodsl/ptodsl/_surface_types.py +++ b/ptodsl/ptodsl/_surface_types.py @@ -286,6 +286,7 @@ class VcvtPartMode: RecipPrecision = _pto.RecipPrecision RsqrtPrecision = _pto.RsqrtPrecision SqrtPrecision = _pto.SqrtPrecision +PrintFormat = _pto.PrintFormat class TensorView: @@ -335,6 +336,7 @@ class Tile: "RecipPrecision", "RsqrtPrecision", "SqrtPrecision", + "PrintFormat", "TensorView", "PartitionTensorView", "Tile", diff --git a/ptodsl/ptodsl/_tile_namespace.py b/ptodsl/ptodsl/_tile_namespace.py index 26cbaeafef..a2c600b340 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/ptodsl/pto.py b/ptodsl/ptodsl/pto.py index e5ea82e48d..5bff50c16a 100644 --- a/ptodsl/ptodsl/pto.py +++ b/ptodsl/ptodsl/pto.py @@ -73,6 +73,7 @@ RecipPrecision, RsqrtPrecision, SqrtPrecision, + PrintFormat, TensorView, PartitionTensorView, Tile, diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index ff9e153d5d..fe5370bb08 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -678,6 +678,34 @@ 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", backend="emitc") +def tile_print_format_emitc_backend_probe(tmp_ptr: pto.ptr(pto.f32, "gm")): + src = pto.alloc_tile(shape=[1, 16], dtype=pto.f32) + tmp_view = pto.make_tensor_view(tmp_ptr, shape=[1, 16], strides=[16, 1]) + tmp = pto.partition_view(tmp_view, offsets=[0, 0], sizes=[1, 16]) + pto.tile.print(src, tmp=tmp, print_format="width8_precision2") + + +@pto.jit(target="a5", backend="emitc") +def tile_print_tmp_format_emitc_backend_probe(tmp_ptr: pto.ptr(pto.f32, "gm")): + src = pto.alloc_tile(shape=[1, 16], dtype=pto.f32) + tmp_view = pto.make_tensor_view(tmp_ptr, shape=[1, 16], strides=[16, 1]) + tmp = pto.partition_view(tmp_view, offsets=[0, 0], sizes=[1, 16]) + pto.tile.print(src, tmp=tmp, print_format="width10_precision6") + + @pto.jit(target="a5") def tile_surface_window_matmul_probe(): src_mat = pto.alloc_tile( @@ -4789,6 +4817,7 @@ 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), @@ -4811,7 +4840,16 @@ def fake_artifacts(py_name, ir_function_name, specialization_key): manifest_path=cache_dir / "manifest.json", ) - def fake_run_ptoas(mlir_path, kernel_object, *, target_arch, insert_sync=None, backend=None, pto_level=None): + def fake_run_ptoas( + mlir_path, + kernel_object, + *, + target_arch, + insert_sync=None, + backend=None, + pto_level=None, + fatobj=False, + ): native_build_observations.append( { "mlir_path": mlir_path, @@ -4820,6 +4858,7 @@ def fake_run_ptoas(mlir_path, kernel_object, *, target_arch, insert_sync=None, b "insert_sync": insert_sync, "backend": backend, "pto_level": pto_level, + "fatobj": fatobj, "mlir_text": mlir_path.read_text(encoding="utf-8"), } ) @@ -4908,6 +4947,11 @@ def fake_link_shared_library(launch_object, kernel_object, shared_library, *, ke observation["pto_level"] == expected_pto_level, f"{label} native build should derive the PTOAS level from the authored mode", ) + expected_fatobj = native_build_runtime._requires_fatobj(compiled.mlir_text()) + expect( + observation["fatobj"] == expected_fatobj, + f"{label} native build should request ptoas fatobj emission only when required by the MLIR", + ) expect( observation["mlir_text"] == compiled.mlir_text(), f"{label} native build should hand the backend-partitioned container MLIR to ptoas unchanged", @@ -4935,6 +4979,7 @@ def fake_run_ptoas_cmd(cmd, *, cwd=None): mlir_path, kernel_object, target_arch="a5", + fatobj=True, ) expect(len(ptoas_cmds) == 1, "native build should issue exactly one ptoas command per kernel container") @@ -4955,6 +5000,10 @@ def fake_run_ptoas_cmd(cmd, *, cwd=None): "--enable-insert-sync" not in ptoas_cmd, "native build should keep the default insert-sync policy unset when _run_ptoas is called directly", ) + expect( + "--fatobj" in ptoas_cmd, + "native build should pass --fatobj when ptoas is responsible for fatobj emission", + ) expect( "--enable-tile-op-expand" in ptoas_cmd and str(mlir_path) in ptoas_cmd and str(kernel_object) in ptoas_cmd, "native build should still pass the shared PTOAS compile inputs and output path", @@ -4970,6 +5019,10 @@ def fake_run_ptoas_cmd(cmd, *, cwd=None): insert_sync=True, ) expect(len(ptoas_cmds) == 1, "native build should issue exactly one ptoas command when insert_sync is forced on") + expect( + "--fatobj" not in ptoas_cmds[0], + "native build should not pass --fatobj unless fatobj emission is explicitly requested", + ) expect( "--enable-insert-sync" in ptoas_cmds[0], "native build should pass --enable-insert-sync when the compiled module explicitly requests it", @@ -5098,6 +5151,38 @@ 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_print_format_text = tile_print_format_emitc_backend_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify( + tile_print_format_text, + "tile print format EmitC backend specialization", + ) + expect( + "printFormat = #pto" in tile_print_format_text, + "pto.tile.print(tmp=..., print_format='width8_precision2') should preserve the print format attribute", + ) + tile_print_tmp_format_text = tile_print_tmp_format_emitc_backend_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify( + tile_print_tmp_format_text, + "tile print tmp format EmitC backend specialization", + ) + expect( + "printFormat = #pto" in tile_print_tmp_format_text, + "pto.tile.print(tmp=..., print_format='width10_precision6') should preserve the print format attribute", + ) 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/python/pto/dialects/pto.py b/python/pto/dialects/pto.py index 161cce5404..87764aa806 100644 --- a/python/pto/dialects/pto.py +++ b/python/pto/dialects/pto.py @@ -105,6 +105,8 @@ def _export_optional_cext_symbol(name): SqrtPrecisionAttr = _pto_mod.SqrtPrecisionAttr FmodPrecision = _pto_mod.FmodPrecision FmodPrecisionAttr = _pto_mod.FmodPrecisionAttr +PrintFormat = _pto_mod.PrintFormat +PrintFormatAttr = _pto_mod.PrintFormatAttr SaturationMode = _pto_mod.SaturationMode SaturationModeAttr = _pto_mod.SaturationModeAttr CmpMode = _pto_mod.CmpMode @@ -284,6 +286,8 @@ def fence_scope_attr_builder(value, context=None): "SqrtPrecisionAttr", "FmodPrecision", "FmodPrecisionAttr", + "PrintFormat", + "PrintFormatAttr", "SaturationMode", "SaturationModeAttr", "CmpMode", diff --git a/test/lit/pto/backend_single_emitc_child_tprint_view.pto b/test/lit/pto/backend_single_emitc_child_tprint_view.pto new file mode 100644 index 0000000000..5a4172590c --- /dev/null +++ b/test/lit/pto/backend_single_emitc_child_tprint_view.pto @@ -0,0 +1,29 @@ +// 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. + +// RUN: ptoas --pto-arch=a5 --enable-tile-op-expand %s -o - 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + module attributes { + pto.backend = "emitc" + } { + func.func @single_emitc_child_tprint_view(%src: !pto.ptr) attributes {pto.entry} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + + %view = pto.make_tensor_view %src, shape = [%c1, %c8], strides = [%c8, %c1] : !pto.tensor_view + %part = pto.partition_view %view, offsets = [%c0, %c0], sizes = [%c1, %c8] : !pto.tensor_view -> !pto.partition_tensor_view<1x8xf32> + pto.tprint ins(%part : !pto.partition_tensor_view<1x8xf32>) + return + } + } +} + +// CHECK-LABEL: AICORE void single_emitc_child_tprint_view( +// CHECK: TPRINT( diff --git a/test/lit/pto/emitc_fatobj_driver_invalid.pto b/test/lit/pto/emitc_fatobj_driver_invalid.pto new file mode 100644 index 0000000000..4bf78ae6c0 --- /dev/null +++ b/test/lit/pto/emitc_fatobj_driver_invalid.pto @@ -0,0 +1,8 @@ +// RUN: not ptoas --pto-arch=a5 --fatobj %s -o - 2>&1 | FileCheck %s --check-prefix=NOOUT +// RUN: not ptoas --pto-arch=a5 --fatobj --emit-pto-ir %s -o %t.o 2>&1 | FileCheck %s --check-prefix=DEBUG + +module { +} + +// NOOUT: Error: --fatobj requires an explicit file path passed with -o +// DEBUG: Error: --fatobj does not support debug IR output flags diff --git a/test/lit/pto/tprint_alloc_tile_no_rebind.pto b/test/lit/pto/tprint_alloc_tile_no_rebind.pto index f219123e55..86a664f120 100644 --- a/test/lit/pto/tprint_alloc_tile_no_rebind.pto +++ b/test/lit/pto/tprint_alloc_tile_no_rebind.pto @@ -1,4 +1,4 @@ -// RUN: ptoas %s | FileCheck %s +// RUN: ptoas --pto-arch=a5 %s | FileCheck %s module { func.func @print_kernel() attributes {pto.entry} { @@ -12,6 +12,16 @@ module { } } +// CHECK: #include +// CHECK: #include +// CHECK: #define conditional PTOAS_ASC_PRINTF_CONDITIONAL +// CHECK: #include "utils/debug/asc_printf.h" +// CHECK: #undef conditional +// CHECK: #include "pto/pto-inst.hpp" +// CHECK: namespace cce { +// CHECK: AICORE inline void printf(const __gm__ char *fmt, Args &&...args) { +// CHECK: __asc_aicore::printf(fmt, args...); +// CHECK: #include "pto/npu/a5/TPrint.hpp" // CHECK-LABEL: __global__ AICORE void print_kernel() { // CHECK: Tile [[TILE_STORAGE:[_A-Za-z][_A-Za-z0-9]*]] // CHECK: Tile [[TILE:[_A-Za-z][_A-Za-z0-9]*]] = [[TILE_STORAGE]]; diff --git a/test/lit/pto/tprint_format_requires_tmp.pto b/test/lit/pto/tprint_format_without_tmp_emitc.pto similarity index 83% rename from test/lit/pto/tprint_format_requires_tmp.pto rename to test/lit/pto/tprint_format_without_tmp_emitc.pto index ee5295c54a..6586976cdd 100644 --- a/test/lit/pto/tprint_format_requires_tmp.pto +++ b/test/lit/pto/tprint_format_without_tmp_emitc.pto @@ -5,10 +5,10 @@ // 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. -// RUN: not ptoas --pto-arch=a3 %s 2>&1 | FileCheck %s +// RUN: not ptoas --pto-arch=a5 %s 2>&1 | FileCheck %s module { - func.func @tprint_format_requires_tmp() attributes {pto.entry} { + func.func @tprint_format_without_tmp() attributes {pto.entry} { %src = pto.alloc_tile : !pto.tile_buf pto.tprint ins(%src : !pto.tile_buf) {printFormat = #pto} @@ -16,4 +16,4 @@ module { } } -// CHECK: error: 'pto.tprint' op expects printFormat only when tmp is present +// CHECK: 'pto.tprint' op expects printFormat only when tmp is present diff --git a/test/lit/pto/tprint_insert_sync_pipe_selection.pto b/test/lit/pto/tprint_insert_sync_pipe_selection.pto new file mode 100644 index 0000000000..59ba76f406 --- /dev/null +++ b/test/lit/pto/tprint_insert_sync_pipe_selection.pto @@ -0,0 +1,32 @@ +// 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. +// RUN: ptoas --pto-arch=a5 --pto-level=level3 --enable-insert-sync %s -o - 2>&1 | FileCheck %s + +module { + func.func @tprint_after_tload(%src: !pto.ptr) attributes {pto.entry} { + %c0_i64 = arith.constant 0 : i64 + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %view = pto.make_tensor_view %src, shape = [%c16, %c16], strides = [%c16, %c1] {layout = #pto.layout}: !pto.tensor_view + %part = pto.partition_view %view, offsets = [%c0, %c0], sizes = [%c16, %c16] : !pto.tensor_view -> !pto.partition_tensor_view<16x16xf32> + %tile = pto.alloc_tile addr = %c0_i64 valid_row = %c16 valid_col = %c16 : !pto.tile_buf + pto.tload ins(%part : !pto.partition_tensor_view<16x16xf32>) + outs(%tile : !pto.tile_buf) + pto.tprint ins(%tile : !pto.tile_buf) + return + } +} + +// CHECK-LABEL: AICORE void tprint_after_tload( +// CHECK: TLOAD( +// CHECK-NEXT: set_flag(PIPE_MTE2, PIPE_S, EVENT_ID0); +// CHECK-NEXT: wait_flag(PIPE_MTE2, PIPE_S, EVENT_ID0); +// CHECK-NOT: set_flag(PIPE_MTE2, PIPE_V +// CHECK-NOT: wait_flag(PIPE_MTE2, PIPE_V +// CHECK: TPRINT( 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..52c44da66e --- /dev/null +++ b/test/tilelib-st/a5/tprint/case.py @@ -0,0 +1,172 @@ +#!/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; the +# simulator does not consistently forward device print payloads to msprof +# stdout, so these ST cases validate the observable data path like the other +# TileLib ST cases: load, print, store, then compare the stored tile. + +from pathlib import Path +import sys + +import numpy as np + +if __package__ in {None, ""}: + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from common import auto_main +from common import assert_close +from ptodsl import pto + + +_NP_TO_PTO = { + np.dtype(np.float32): pto.f32, + np.dtype(np.int32): pto.i32, + np.dtype(np.uint32): pto.ui32, + np.dtype(np.int8): pto.i8, + np.dtype(np.uint8): pto.ui8, +} + + +CASE_SPECS = [ + { + "name": "tprint_float_formats_default", + "np_dtype": np.float32, + "shape": (1, 8), + "values": (1.234567, -3.456789), + }, + { + "name": "tprint_float_formats_precision2", + "np_dtype": np.float32, + "shape": (1, 8), + "values": (1.234567, -3.456789), + "use_tmp": True, + "print_format": "width8_precision2", + }, + { + "name": "tprint_float_formats_precision6", + "np_dtype": np.float32, + "shape": (1, 8), + "values": (1.234567, -3.456789), + "use_tmp": True, + "print_format": "width10_precision6", + }, + { + "name": "tprint_signed_int_formats", + "np_dtype": np.int32, + "shape": (1, 8), + "values": (42, -17, 1024, -9999), + }, + { + "name": "tprint_unsigned_int_formats", + "np_dtype": np.uint32, + "shape": (1, 8), + "values": (0, 17, 65535, 123456), + }, + { + "name": "tprint_int8_printed_as_numbers", + "np_dtype": np.int8, + "shape": (1, 32), + "values": (-12, 65), + }, + { + "name": "tprint_uint8_printed_as_numbers", + "np_dtype": np.uint8, + "shape": (1, 32), + "values": (255, 127), + }, + { + "name": "tprint_tile_shape_header", + "np_dtype": np.float32, + "shape": (2, 8), + "values": (1.0,), + }, + { + "name": "tprint_overload_with_tmp", + "np_dtype": np.float32, + "shape": (1, 8), + "values": (3.141592,), + "use_tmp": True, + "print_format": "width10_precision6", + }, +] + + +def _make_inputs(spec): + data = np.zeros(spec["shape"], dtype=spec["np_dtype"]) + flat = data.reshape(-1) + for idx, value in enumerate(spec["values"]): + flat[idx] = value + out = np.zeros(spec["shape"], dtype=spec["np_dtype"]) + if spec.get("use_tmp", False): + tmp = np.zeros(spec["shape"], dtype=spec["np_dtype"]) + return [data, tmp, out] + return [data, out] + + +def _make_kernel(spec): + dtype = _NP_TO_PTO[np.dtype(spec["np_dtype"])] + rows, cols = spec["shape"] + use_tmp = spec.get("use_tmp", False) + print_format = spec.get("print_format") + kernel_name = spec["name"] + + if use_tmp: + @pto.jit(name=kernel_name, target="a5", backend="emitc") + def _kernel( + src_ptr: pto.ptr(dtype, "gm"), + tmp_ptr: pto.ptr(dtype, "gm"), + out_ptr: pto.ptr(dtype, "gm"), + ): + src_view = pto.make_tensor_view(src_ptr, shape=[rows, cols], strides=[cols, 1]) + tmp_view = pto.make_tensor_view(tmp_ptr, shape=[rows, cols], strides=[cols, 1]) + out_view = pto.make_tensor_view(out_ptr, shape=[rows, cols], strides=[cols, 1]) + tmp = pto.partition_view(tmp_view, offsets=[0, 0], sizes=[rows, cols]) + tile = pto.alloc_tile(shape=[rows, cols], dtype=dtype) + + pto.tile.load(src_view, tile) + pto.tile.print(tile, tmp=tmp, print_format=print_format) + pto.tile.store(tile, out_view) + else: + @pto.jit(name=kernel_name, target="a5", backend="emitc") + def _kernel(src_ptr: pto.ptr(dtype, "gm"), out_ptr: pto.ptr(dtype, "gm")): + src_view = pto.make_tensor_view(src_ptr, shape=[rows, cols], strides=[cols, 1]) + out_view = pto.make_tensor_view(out_ptr, shape=[rows, cols], strides=[cols, 1]) + tile = pto.alloc_tile(shape=[rows, cols], dtype=dtype) + + pto.tile.load(src_view, tile) + pto.tile.print(tile, print_format=print_format) + pto.tile.store(tile, out_view) + + return _kernel + + +def _make_case(spec): + inputs = _make_inputs(spec) + return inputs, inputs[0] + + +def _check_case(device_inputs, golden): + actual = device_inputs[-1].cpu().numpy() + assert_close(actual, golden, rtol=1e-6, atol=1e-6) + + +CASES = [] +for _spec in CASE_SPECS: + CASES.append( + { + "name": _spec["name"], + "kernel": _make_kernel(_spec), + "make_case": lambda _spec=_spec: _make_case(_spec), + "check": _check_case, + } + ) + + +auto_main(globals()) diff --git a/tools/ptoas/ObjectEmission.cpp b/tools/ptoas/ObjectEmission.cpp index 806b0cfc9f..6cca39a09d 100644 --- a/tools/ptoas/ObjectEmission.cpp +++ b/tools/ptoas/ObjectEmission.cpp @@ -241,6 +241,16 @@ discoverCppIncludeDirs(llvm::StringRef ascendHome, addPTOISAIncludeDirs(includeDirs, ptoIsaPath); addExistingIncludeDir(includeDirs, joinPath(ascendHome, "include")); + addExistingIncludeDir(includeDirs, joinPath(ascendHome, "aarch64-linux/asc")); + addExistingIncludeDir(includeDirs, + joinPath(ascendHome, "aarch64-linux/asc/include")); + addExistingIncludeDir(includeDirs, joinPath(ascendHome, "x86_64-linux/asc")); + addExistingIncludeDir(includeDirs, + joinPath(ascendHome, "x86_64-linux/asc/include")); + addExistingIncludeDir(includeDirs, joinPath(ascendHome, "pkg_inc")); + addExistingIncludeDir(includeDirs, joinPath(ascendHome, "pkg_inc/profiling")); + addExistingIncludeDir(includeDirs, + joinPath(ascendHome, "pkg_inc/runtime/runtime")); std::string driverPath = getEnvPath("ASCEND_DRIVER_PATH").value_or("/usr/local/Ascend/driver"); addExistingIncludeDir(includeDirs, joinPath(driverPath, "kernel/inc")); @@ -571,7 +581,7 @@ static bool compileCppDeviceSourceToObject( static bool compileCppDeviceSourceToFatobj( llvm::StringRef cppPath, llvm::StringRef outObjPath, - const mlir::pto::CANNToolchain &toolchain, + llvm::StringRef targetCPU, const mlir::pto::CANNToolchain &toolchain, llvm::StringRef stderrPath, llvm::raw_ostream &diagOS) { llvm::SmallVector args = { toolchain.bishengPath, @@ -591,7 +601,7 @@ static bool compileCppDeviceSourceToFatobj( "-cce-aicore-addr-transform", "-mllvm", "-cce-aicore-dcci-insert-for-scalar=false", - "--cce-aicore-arch=dav-c310", + std::string("--cce-aicore-arch=") + targetCPU.str(), "-DREGISTER_BASE", "-std=c++17", "-O2", @@ -752,6 +762,7 @@ static bool mergeDeviceObjects(llvm::ArrayRef deviceObjPaths, static bool linkFatobjFiles(llvm::ArrayRef fatobjPaths, llvm::StringRef outObjPath, + llvm::StringRef targetCPU, const mlir::pto::CANNToolchain &toolchain, llvm::StringRef stderrPath, llvm::raw_ostream &diagOS) { @@ -761,7 +772,7 @@ static bool linkFatobjFiles(llvm::ArrayRef fatobjPaths, llvm::SmallVector args = { toolchain.bishengPath, "--cce-fatobj-link", - "--cce-aicore-arch=dav-c310", + std::string("--cce-aicore-arch=") + targetCPU.str(), "-r", "-o", outObjPath.str(), @@ -935,28 +946,29 @@ mlir::LogicalResult mlir::pto::emitCppCubeDeviceObject( mlir::LogicalResult mlir::pto::emitCppFatobj( llvm::StringRef cppSource, llvm::StringRef cppPath, - llvm::StringRef outObjPath, const CANNToolchain &toolchain, - llvm::StringRef stderrPath, llvm::raw_ostream &diagOS) { + llvm::StringRef outObjPath, llvm::StringRef targetCPU, + const CANNToolchain &toolchain, llvm::StringRef stderrPath, + llvm::raw_ostream &diagOS) { if (failed(writeCppSource(cppSource, cppPath, diagOS))) return failure(); - return compileCppDeviceSourceToFatobj(cppPath, outObjPath, toolchain, - stderrPath, diagOS) + return compileCppDeviceSourceToFatobj(cppPath, outObjPath, targetCPU, + toolchain, stderrPath, diagOS) ? success() : failure(); } mlir::LogicalResult mlir::pto::emitFatobjCCE( llvm::StringRef cppSource, llvm::StringRef outputPath, - const CANNToolchain &toolchain, TempFileRegistry &tempFiles, - llvm::raw_ostream &diagOS) { + llvm::StringRef targetCPU, const CANNToolchain &toolchain, + TempFileRegistry &tempFiles, llvm::raw_ostream &diagOS) { std::string cppPath; std::string stderrPath; if (failed(tempFiles.create("ptoas-emitc", ".cpp", cppPath, diagOS)) || failed(tempFiles.create("ptoas-emitc-fatobj", ".log", stderrPath, diagOS))) return failure(); - return emitCppFatobj(cppSource, cppPath, outputPath, toolchain, stderrPath, - diagOS); + return emitCppFatobj(cppSource, cppPath, outputPath, targetCPU, toolchain, + stderrPath, diagOS); } static bool isVPTOKernelABISymbol(llvm::StringRef name) { @@ -1043,8 +1055,9 @@ mlir::LogicalResult mlir::pto::emitVPTOCubeDeviceObject( mlir::LogicalResult mlir::pto::emitFatobjLLVM( llvm::Module *cubeModule, llvm::Module *vectorModule, llvm::StringRef stubSource, llvm::StringRef outputPath, - llvm::StringRef moduleId, const CANNToolchain &toolchain, - TempFileRegistry &tempFiles, VFSIMTSizeFixMode vfsimtSizeFixMode, + llvm::StringRef moduleId, llvm::StringRef targetCPU, + const CANNToolchain &toolchain, TempFileRegistry &tempFiles, + VFSIMTSizeFixMode vfsimtSizeFixMode, llvm::raw_ostream &diagOS) { if (!cubeModule && !vectorModule) { diagOS << "Error: VPTO fatobj emission requires at least one LLVM module.\n"; @@ -1064,7 +1077,6 @@ mlir::LogicalResult mlir::pto::emitFatobjLLVM( if (!artifacts.mergeDeviceObjects(toolchain, diagOS)) return failure(); - constexpr llvm::StringLiteral targetCPU = "dav-c310"; if (!artifacts.compileHostStubToFatobj(toolchain, moduleId, targetCPU, outputPath, diagOS)) return failure(); @@ -1084,9 +1096,8 @@ mlir::LogicalResult mlir::pto::mergeDeviceObjects( mlir::LogicalResult mlir::pto::compileStubToFatobj( llvm::StringRef stubPath, llvm::StringRef deviceObjPath, llvm::StringRef outputPath, llvm::StringRef moduleId, - const CANNToolchain &toolchain, llvm::StringRef stderrPath, - llvm::raw_ostream &diagOS) { - constexpr llvm::StringLiteral targetCPU = "dav-c310"; + llvm::StringRef targetCPU, const CANNToolchain &toolchain, + llvm::StringRef stderrPath, llvm::raw_ostream &diagOS) { return compileHostStubToObject(stubPath, outputPath, moduleId, targetCPU, toolchain, deviceObjPath, stderrPath, diagOS) @@ -1096,9 +1107,10 @@ mlir::LogicalResult mlir::pto::compileStubToFatobj( mlir::LogicalResult mlir::pto::linkFatobjs( llvm::ArrayRef fatobjPaths, llvm::StringRef outputPath, - const CANNToolchain &toolchain, llvm::StringRef stderrPath, - llvm::raw_ostream &diagOS) { - return linkFatobjFiles(fatobjPaths, outputPath, toolchain, stderrPath, diagOS) + llvm::StringRef targetCPU, const CANNToolchain &toolchain, + llvm::StringRef stderrPath, llvm::raw_ostream &diagOS) { + return linkFatobjFiles(fatobjPaths, outputPath, targetCPU, toolchain, + stderrPath, diagOS) ? success() : failure(); } diff --git a/tools/ptoas/ObjectEmission.h b/tools/ptoas/ObjectEmission.h index 11232c5873..d44626a012 100644 --- a/tools/ptoas/ObjectEmission.h +++ b/tools/ptoas/ObjectEmission.h @@ -106,12 +106,14 @@ LogicalResult emitCppCubeDeviceObject( LogicalResult emitCppFatobj(llvm::StringRef cppSource, llvm::StringRef cppPath, llvm::StringRef outObjPath, + llvm::StringRef targetCPU, const CANNToolchain &toolchain, llvm::StringRef stderrPath, llvm::raw_ostream &diagOS); LogicalResult emitFatobjCCE(llvm::StringRef cppSource, llvm::StringRef outputPath, + llvm::StringRef targetCPU, const CANNToolchain &toolchain, TempFileRegistry &tempFiles, llvm::raw_ostream &diagOS); @@ -129,8 +131,9 @@ LogicalResult emitVPTOCubeDeviceObject( LogicalResult emitFatobjLLVM( llvm::Module *cubeModule, llvm::Module *vectorModule, llvm::StringRef stubSource, llvm::StringRef outputPath, - llvm::StringRef moduleId, const CANNToolchain &toolchain, - TempFileRegistry &tempFiles, VFSIMTSizeFixMode vfsimtSizeFixMode, + llvm::StringRef moduleId, llvm::StringRef targetCPU, + const CANNToolchain &toolchain, TempFileRegistry &tempFiles, + VFSIMTSizeFixMode vfsimtSizeFixMode, llvm::raw_ostream &diagOS); LogicalResult mergeDeviceObjects(llvm::ArrayRef deviceObjPaths, @@ -142,11 +145,12 @@ LogicalResult mergeDeviceObjects(llvm::ArrayRef deviceObjPaths, LogicalResult compileStubToFatobj( llvm::StringRef stubPath, llvm::StringRef deviceObjPath, llvm::StringRef outputPath, llvm::StringRef moduleId, - const CANNToolchain &toolchain, llvm::StringRef stderrPath, - llvm::raw_ostream &diagOS); + llvm::StringRef targetCPU, const CANNToolchain &toolchain, + llvm::StringRef stderrPath, llvm::raw_ostream &diagOS); LogicalResult linkFatobjs(llvm::ArrayRef fatobjPaths, llvm::StringRef outputPath, + llvm::StringRef targetCPU, const CANNToolchain &toolchain, llvm::StringRef stderrPath, llvm::raw_ostream &diagOS); diff --git a/tools/ptoas/driver.cpp b/tools/ptoas/driver.cpp index 4c44d980ce..28b456c70d 100644 --- a/tools/ptoas/driver.cpp +++ b/tools/ptoas/driver.cpp @@ -60,6 +60,11 @@ static llvm::cl::opt outputFilename("o", llvm::cl::desc("Output filename"), llvm::cl::value_desc("filename"), llvm::cl::init("-")); +llvm::cl::opt mlir::pto::emitFatobj( + "fatobj", + llvm::cl::desc("Compile backend output to a CCE fat object"), + llvm::cl::init(false)); + static void printPTOASVersion(llvm::raw_ostream &os) { os << "ptoas " << PTOAS_RELEASE_VERSION << "\n"; } @@ -345,6 +350,117 @@ static SmallVector collectDirectCalleeNames(func::FuncOp funcOp) { return names; } +static llvm::StringSet<> collectSiblingDirectCalleeNames(ModuleOp outer, + ModuleOp targetChild) { + llvm::StringSet<> names; + for (ModuleOp child : outer.getOps()) { + if (child == targetChild) + continue; + for (StringRef calleeName : collectDirectCalleeNames(child)) + names.insert(calleeName); + } + return names; +} + +static std::optional +getChildKernelKind(ModuleOp module) { + auto kernelKindAttr = + module->getAttrOfType( + mlir::pto::FunctionKernelKindAttr::name); + if (!kernelKindAttr) + return std::nullopt; + return kernelKindAttr.getKernelKind(); +} + +static std::string computeFatobjTargetCPU( + llvm::StringRef ptoArch, + std::optional kernelKind = std::nullopt) { + const bool isA5 = normalizePTOASArch(ptoArch) == "a5"; + llvm::StringRef base = isA5 ? "dav-c310" : "dav-c220"; + if (!kernelKind) + return base.str(); + switch (*kernelKind) { + case mlir::pto::FunctionKernelKind::Vector: + return (base + "-vec").str(); + case mlir::pto::FunctionKernelKind::Cube: + return (base + "-cube").str(); + } + return base.str(); +} + +static std::string computeEmitCFatobjTargetCPU( + llvm::StringRef ptoArch, + std::optional kernelKind, + llvm::StringRef cppOutput) { + if (kernelKind) + return computeFatobjTargetCPU(ptoArch, kernelKind); + if (cppOutput.contains("__DAV_CUBE__")) + return computeFatobjTargetCPU(ptoArch, + mlir::pto::FunctionKernelKind::Cube); + return computeFatobjTargetCPU(ptoArch, mlir::pto::FunctionKernelKind::Vector); +} + +static bool hasVPTOPublicABISuffix(llvm::StringRef symbolName) { + return symbolName.ends_with(".vector") || symbolName.ends_with(".cube") || + symbolName.ends_with("_mix_aiv") || symbolName.ends_with("_mix_aic"); +} + +static FailureOr +getCrossChildCallSymbolName(func::FuncOp siblingSource, + mlir::pto::PTOBackend consumerBackend, + mlir::pto::PTOBackend defaultBackend, + const mlir::pto::CANNToolchain *toolchain) { + ModuleOp providerChild = siblingSource->getParentOfType(); + std::optional providerBackend; + if (failed(parseDriverBackendAttr(providerChild.getOperation(), + providerBackend))) + return failure(); + + if (consumerBackend == mlir::pto::PTOBackend::EmitC) + return siblingSource.getSymName().str(); + + if (providerBackend.value_or(defaultBackend) != mlir::pto::PTOBackend::VPTO || + !toolchain) + return siblingSource.getSymName().str(); + + std::optional kernelKind = + getChildKernelKind(providerChild); + if (!kernelKind) { + providerChild.emitError( + "mixed-backend child assembly cannot form a VPTO ABI symbol for cross-child function reference '@") + << siblingSource.getSymName() << "' without " + << mlir::pto::FunctionKernelKindAttr::name; + return failure(); + } + + mlir::pto::ObjectEmissionDeviceTarget target; + if (*kernelKind == mlir::pto::FunctionKernelKind::Vector) { + target = mlir::pto::ObjectEmissionDeviceTarget::Vector; + } else if (*kernelKind == mlir::pto::FunctionKernelKind::Cube) { + target = mlir::pto::ObjectEmissionDeviceTarget::Cube; + } else { + providerChild.emitError( + "mixed-backend child assembly cannot form a VPTO ABI symbol for unsupported cross-child function kind"); + return failure(); + } + std::string symbolName = siblingSource.getSymName().str(); + if (hasVPTOPublicABISuffix(symbolName)) + return symbolName; + return (llvm::StringRef(symbolName) + toolchain->vptoPublicABISuffix(target)) + .str(); +} + +static void rewriteDirectCallCallees(ModuleOp module, StringRef oldName, + StringRef newName) { + if (oldName == newName) + return; + module.walk([&](func::CallOp callOp) { + if (callOp.getCalleeAttr().getLeafReference() != oldName) + return; + callOp.setCalleeAttr(FlatSymbolRefAttr::get(module.getContext(), newName)); + }); +} + static void copyModuleAttrsToJobModule(ModuleOp source, ModuleOp jobModule) { for (NamedAttribute attr : source->getAttrs()) { StringRef attrName = attr.getName().getValue(); @@ -517,7 +633,17 @@ verifyInChildLogicalWrapperAmbiguity(ModuleOp targetChild, } static FailureOr> -buildBackendChildCompileUnit(ModuleOp outer, ModuleOp targetChild) { +buildBackendChildCompileUnit(ModuleOp outer, ModuleOp targetChild, + mlir::pto::PTOBackend defaultBackend = + mlir::pto::PTOBackend::EmitC, + const mlir::pto::CANNToolchain *toolchain = + nullptr) { + std::optional targetBackend; + if (failed(parseDriverBackendAttr(targetChild.getOperation(), targetBackend))) + return failure(); + mlir::pto::PTOBackend effectiveTargetBackend = + targetBackend.value_or(defaultBackend); + ModuleOp jobModule = ModuleOp::create(outer.getLoc()); copyModuleAttrsToJobModule(outer, jobModule); copyModuleAttrsToJobModule(targetChild, jobModule); @@ -526,10 +652,22 @@ buildBackendChildCompileUnit(ModuleOp outer, ModuleOp targetChild) { jobModule.push_back(op.clone()); } + llvm::StringSet<> siblingDirectCallees = + collectSiblingDirectCalleeNames(outer, targetChild); + for (func::FuncOp funcOp : jobModule.getOps()) { + if (funcOp.isDeclaration() || funcOp.isPrivate()) + continue; + if (!siblingDirectCallees.contains(funcOp.getSymName())) + continue; + mlir::pto::setExternalArtifactVisibility(funcOp, true); + } + SmallVector directCalleeNames = collectDirectCalleeNames(targetChild); for (StringRef calleeName : directCalleeNames) { - if (findFunctionByLogicalName(jobModule, calleeName)) + func::FuncOp localFunc = findFunctionByLogicalName(jobModule, calleeName); + if (localFunc && !localFunc.isDeclaration()) continue; + FailureOr siblingSourceOr = findSiblingSourceFunction(outer, targetChild, calleeName, /*allowLogicalNameMatch=*/false, @@ -537,6 +675,8 @@ buildBackendChildCompileUnit(ModuleOp outer, ModuleOp targetChild) { if (failed(siblingSourceOr)) return failure(); func::FuncOp siblingSource = *siblingSourceOr; + if (!siblingSource && localFunc) + continue; if (!siblingSource) { targetChild.emitError( "mixed-backend child assembly does not yet support unresolved cross-child function reference '@") @@ -544,8 +684,27 @@ buildBackendChildCompileUnit(ModuleOp outer, ModuleOp targetChild) { << "'; each cross-child func.call must resolve to one sibling public func.func"; return failure(); } - cloneFunctionDeclarationIntoModule(jobModule, siblingSource, calleeName, - "private"); + FailureOr importNameOr = + getCrossChildCallSymbolName(siblingSource, effectiveTargetBackend, + defaultBackend, toolchain); + if (failed(importNameOr)) + return failure(); + StringRef importName = *importNameOr; + if (func::FuncOp existing = findFunctionBySymbolName(jobModule, importName); + existing && existing != localFunc) { + targetChild.emitError( + "mixed-backend child assembly cannot import cross-child function '@") + << calleeName << "' as '@" << importName + << "' because that symbol already exists in the child compile unit"; + return failure(); + } + if (localFunc) { + localFunc.setSymName(importName); + } else { + cloneFunctionDeclarationIntoModule(jobModule, siblingSource, importName, + "private"); + } + rewriteDirectCallCallees(jobModule, calleeName, importName); } SmallVector importedPeerNames = collectImportedPeerNames(targetChild); @@ -674,6 +833,9 @@ static LogicalResult emitVPTOLLVMFatobj( const mlir::pto::PTOASCompileResult &jobResult, mlir::pto::PTOASContext &context, llvm::StringRef moduleId, llvm::StringRef outputPath); +static LogicalResult emitEmitCFatobj( + mlir::pto::PTOASCompileResult &jobResult, + mlir::pto::PTOASContext &context, llvm::StringRef targetCPU); mlir::pto::PTOASContext::PTOASContext(DialectRegistry ®istry, llvm::StringRef outputPath, int argc, @@ -766,6 +928,13 @@ mlir::pto::PTOASContext::initializeToolchain(llvm::raw_ostream &diagOS) { return success(); } +const mlir::pto::CANNToolchain * +mlir::pto::PTOASContext::peekToolchain() const { + if (!toolchain) + return nullptr; + return &*toolchain; +} + const mlir::pto::CANNToolchain * mlir::pto::PTOASContext::getToolchain(llvm::raw_ostream &diagOS) const { if (!toolchain) { @@ -811,10 +980,12 @@ class EmitCBackendJob { : module(module), result(result) {} LogicalResult run(PTOASContext &context); + llvm::StringRef getFatobjTargetCPU() const { return fatobjTargetCPU; } private: OwningOpRef &module; mlir::pto::PTOASCompileResult &result; + std::string fatobjTargetCPU; }; class VPTOBackendJob { @@ -840,9 +1011,10 @@ class EmitCBackendChildJob final : public BackendChildJob { public: EmitCBackendChildJob(OwningOpRef &&module, std::string summary, - SmallVectorImpl &fatobjPaths) + SmallVectorImpl &fatobjPaths, + bool forceBaseTargetCPU) : module(std::move(module)), summary(std::move(summary)), - fatobjPaths(fatobjPaths) {} + fatobjPaths(fatobjPaths), forceBaseTargetCPU(forceBaseTargetCPU) {} LogicalResult run(PTOASContext &context) override { ModuleOp op = module.get(); @@ -870,7 +1042,13 @@ class EmitCBackendChildJob final : public BackendChildJob { if (!toolchain) return failure(); if (failed(mlir::pto::emitFatobjCCE( - jobResult.textOutput, fatobjPath, *toolchain, + jobResult.textOutput, fatobjPath, + forceBaseTargetCPU + ? computeFatobjTargetCPU(context.getArch()) + : computeEmitCFatobjTargetCPU(context.getArch(), + getChildKernelKind(op), + jobResult.textOutput), + *toolchain, context.getTempFiles(), llvm::errs()))) { dumpFailedMixedChildCompileUnit("emitc", summary, op); return failure(); @@ -884,6 +1062,7 @@ class EmitCBackendChildJob final : public BackendChildJob { OwningOpRef module; std::string summary; SmallVectorImpl &fatobjPaths; + bool forceBaseTargetCPU; }; class VPTOBackendChildJob final : public BackendChildJob { @@ -953,8 +1132,10 @@ class FatobjLinkJob { context.getToolchain(llvm::errs()); if (!toolchain) return failure(); - return mlir::pto::linkFatobjs(fatobjPaths, context.getOutputPath(), - *toolchain, stderrPath, llvm::errs()); + return mlir::pto::linkFatobjs( + fatobjPaths, context.getOutputPath(), + computeFatobjTargetCPU(context.getArch()), *toolchain, stderrPath, + llvm::errs()); } private: @@ -968,6 +1149,9 @@ LogicalResult EmitCBackendJob::run(PTOASContext &context) { op->setAttr("pto.backend", StringAttr::get(op.getContext(), "emitc")); SmallVector children(op.getOps()); + // PTODSL uses a backend-partitioned outer module for both explicit helper + // children and ordinary EmitC entry kernels. Compile the single child in the + // same normalized job shape so frontend view ops are lowered before EmitC. if (!isUserVisibleIROutputRequested() && children.size() == 1 && isBackendPartitionedContainer(op)) { FailureOr> jobModuleOr = @@ -989,6 +1173,9 @@ LogicalResult EmitCBackendJob::run(PTOASContext &context) { llvm::errs() << "Error: EmitC backend job produced non-text output.\n"; return failure(); } + fatobjTargetCPU = computeEmitCFatobjTargetCPU( + context.getArch(), getChildKernelKind(compileUnit->get()), + result.textOutput); return success(); } @@ -1060,25 +1247,62 @@ static LogicalResult emitVPTOLLVMFatobj( if (failed(mlir::pto::emitFatobjLLVM( jobResult.vptoCubeModule.module.get(), jobResult.vptoVectorModule.module.get(), stubSource, - outputPath, moduleId, *toolchain, context.getTempFiles(), - context.getVFSIMTSizeFixMode(), llvm::errs()))) + outputPath, moduleId, computeFatobjTargetCPU(context.getArch()), + *toolchain, context.getTempFiles(), context.getVFSIMTSizeFixMode(), + llvm::errs()))) return failure(); return success(); } +static LogicalResult emitEmitCFatobj(mlir::pto::PTOASCompileResult &jobResult, + PTOASContext &context, + llvm::StringRef targetCPU) { + if (jobResult.kind != mlir::pto::PTOASCompileResultKind::Text) { + llvm::errs() << "Error: EmitC fatobj mode expected C++ text output.\n"; + return failure(); + } + if (context.getOutputPath().empty() || context.getOutputPath() == "-") { + llvm::errs() << "Error: EmitC fatobj mode requires an explicit file path " + "passed with -o.\n"; + return failure(); + } + + const mlir::pto::CANNToolchain *toolchain = + context.getToolchain(llvm::errs()); + if (!toolchain) + return failure(); + if (failed(mlir::pto::emitFatobjCCE(jobResult.textOutput, + context.getOutputPath(), targetCPU, + *toolchain, + context.getTempFiles(), llvm::errs()))) + return failure(); + + jobResult.reset(); + jobResult.kind = mlir::pto::PTOASCompileResultKind::MixedObject; + return success(); +} + static LogicalResult collectChildJobs( ModuleOp module, mlir::pto::PTOBackend defaultBackend, - bool cliBackendOverride, PTOASContext &context, SmallVectorImpl &fatobjPaths, SmallVectorImpl> &backendJobs) { + const mlir::pto::CANNToolchain *toolchain = context.peekToolchain(); SmallVector children(module.getOps()); + bool hasVPTOChild = false; + for (ModuleOp child : children) { + std::optional childBackend; + if (failed(parseDriverBackendAttr(child.getOperation(), childBackend))) + return failure(); + if (childBackend.value_or(defaultBackend) == mlir::pto::PTOBackend::VPTO) + hasVPTOChild = true; + } for (ModuleOp child : children) { std::optional childBackend; if (failed(parseDriverBackendAttr(child.getOperation(), childBackend))) return failure(); FailureOr> jobModuleOr = - buildBackendChildCompileUnit(module, child); + buildBackendChildCompileUnit(module, child, defaultBackend, toolchain); if (failed(jobModuleOr)) return failure(); OwningOpRef jobModule = std::move(*jobModuleOr); @@ -1089,15 +1313,15 @@ static LogicalResult collectChildJobs( } std::string summary = summarizeMixedChildModule(jobModule.get()); mlir::pto::PTOBackend effectiveBackend = - cliBackendOverride ? defaultBackend - : childBackend.value_or(defaultBackend); + childBackend.value_or(defaultBackend); if (effectiveBackend == mlir::pto::PTOBackend::VPTO) backendJobs.push_back(std::make_unique( std::move(jobModule), std::move(summary), context.allocModuleId(), fatobjPaths)); else backendJobs.push_back(std::make_unique( - std::move(jobModule), std::move(summary), fatobjPaths)); + std::move(jobModule), std::move(summary), fatobjPaths, + /*forceBaseTargetCPU=*/hasVPTOChild)); } return success(); } @@ -1178,11 +1402,23 @@ static LogicalResult buildBackendInfo(ModuleOp module, bool cliBackendSpecified, backendInfo.singleBackend))) return failure(); + if (mlir::pto::emitFatobj && isUserVisibleIROutputRequested()) { + llvm::errs() << "Error: --fatobj does not support debug IR output flags.\n"; + return failure(); + } + if (mlir::pto::emitFatobj && + (outputFilename.empty() || outputFilename == "-")) { + llvm::errs() << "Error: --fatobj requires an explicit file path passed " + "with -o.\n"; + return failure(); + } + if (backendInfo.singleBackend) { backendInfo.requiresToolchain = - *backendInfo.singleBackend == mlir::pto::PTOBackend::VPTO && - !mlir::pto::emitMlirIR && !mlir::pto::emitVPTO && - !mlir::pto::emitVPTOLLVMDialect; + mlir::pto::emitFatobj || + (*backendInfo.singleBackend == mlir::pto::PTOBackend::VPTO && + !mlir::pto::emitMlirIR && !mlir::pto::emitVPTO && + !mlir::pto::emitVPTOLLVMDialect); return success(); } @@ -1210,7 +1446,11 @@ static LogicalResult runPTOASJobs(OwningOpRef &module, if (backendInfo.singleBackend) { if (*backendInfo.singleBackend == mlir::pto::PTOBackend::EmitC) { EmitCBackendJob singleJob(module, result); - return singleJob.run(context); + if (failed(singleJob.run(context))) + return failure(); + if (mlir::pto::emitFatobj) + return emitEmitCFatobj(result, context, singleJob.getFatobjTargetCPU()); + return success(); } VPTOBackendJob singleJob(module, result); return singleJob.run(context); @@ -1219,7 +1459,6 @@ static LogicalResult runPTOASJobs(OwningOpRef &module, SmallVector, 4> backendJobs; SmallVector fatobjPaths; if (failed(collectChildJobs(module.get(), backendInfo.defaultBackend, - backendInfo.cliBackendOverride, context, fatobjPaths, backendJobs))) return failure(); diff --git a/tools/ptoas/ptoas.h b/tools/ptoas/ptoas.h index cf5b369ef7..42037de658 100644 --- a/tools/ptoas/ptoas.h +++ b/tools/ptoas/ptoas.h @@ -35,6 +35,7 @@ extern llvm::cl::opt ptoTargetArch; extern llvm::cl::opt ptoBackend; extern llvm::cl::opt emitVPTO; extern llvm::cl::opt emitVPTOLLVMDialect; +extern llvm::cl::opt emitFatobj; extern llvm::cl::opt ptoPrintSeamIR; extern llvm::cl::opt ptoSeamIRFile; extern llvm::cl::opt cannOutputVersion; @@ -85,6 +86,7 @@ class PTOASContext { llvm::StringRef getOutputPath() const; std::string allocModuleId(); + const CANNToolchain *peekToolchain() const; const CANNToolchain *getToolchain(llvm::raw_ostream &diagOS) const; CANNVersion getCANNVersionOrDefault() const;