From a4aa9009e47c07faf3f2082f205a1afe110acda8 Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Thu, 28 May 2026 23:17:29 +0800 Subject: [PATCH 1/2] feat: first version of driver --- .../module-backend-attr-and-mixed-fatobj.md | 553 ++++ lib/PTO/Transforms/PTOToEmitC.cpp | 8 +- test/lit/vpto/backend_attr_invalid.pto | 21 + test/lit/vpto/backend_attr_selects_vpto.pto | 26 + ...backend_child_attr_inherits_outer_arch.pto | 29 + .../vpto/backend_child_attr_selects_vpto.pto | 29 + test/lit/vpto/backend_cli_override.pto | 23 + .../lit/vpto/backend_cli_override_to_vpto.pto | 24 + .../backend_mixed_child_attrs_rejected.pto | 28 + ...nd_mixed_child_missing_backend_invalid.pto | 34 + ...end_mixed_duplicate_abi_export_invalid.pto | 37 + ...d_mixed_external_import_aicore_invalid.pto | 29 + .../backend_mixed_external_import_missing.pto | 32 + ...backend_mixed_external_import_resolves.pto | 36 + ...xed_external_import_signature_mismatch.pto | 29 + .../backend_mixed_requires_output_file.pto | 28 + ...end_mixed_vpto_stub_signature_conflict.pto | 37 + .../vpto/backend_public_abi_export_suffix.pto | 38 + ...end_reserved_aicore_abi_suffix_invalid.pto | 21 + ...end_reserved_public_abi_suffix_invalid.pto | 21 + ...vpto_child_missing_kernel_kind_invalid.pto | 19 + .../backend/mixed-external-vadd/compare.py | 40 + .../backend/mixed-external-vadd/golden.py | 43 + .../backend/mixed-external-vadd/kernel.pto | 93 + .../backend/mixed-external-vadd/launch.cpp | 25 + .../backend/mixed-external-vadd/main.cpp | 118 + .../backend/mixed-external-vadd/ptoas.flags | 1 + test/vpto/scripts/run_host_vpto_validation.sh | 9 +- tools/ptoas/CMakeLists.txt | 3 +- ...OFatobjEmission.cpp => ObjectEmission.cpp} | 376 ++- tools/ptoas/ObjectEmission.h | 121 + tools/ptoas/PTOASDriver.cpp | 2726 +++++++++++++++++ .../{VPTOFatobjEmission.h => PTOASDriver.h} | 19 +- tools/ptoas/VPTOHostStubEmission.cpp | 60 +- tools/ptoas/VPTOHostStubEmission.h | 4 + tools/ptoas/ptoas.cpp | 1673 +--------- 36 files changed, 4685 insertions(+), 1728 deletions(-) create mode 100644 docs/designs/module-backend-attr-and-mixed-fatobj.md create mode 100644 test/lit/vpto/backend_attr_invalid.pto create mode 100644 test/lit/vpto/backend_attr_selects_vpto.pto create mode 100644 test/lit/vpto/backend_child_attr_inherits_outer_arch.pto create mode 100644 test/lit/vpto/backend_child_attr_selects_vpto.pto create mode 100644 test/lit/vpto/backend_cli_override.pto create mode 100644 test/lit/vpto/backend_cli_override_to_vpto.pto create mode 100644 test/lit/vpto/backend_mixed_child_attrs_rejected.pto create mode 100644 test/lit/vpto/backend_mixed_child_missing_backend_invalid.pto create mode 100644 test/lit/vpto/backend_mixed_duplicate_abi_export_invalid.pto create mode 100644 test/lit/vpto/backend_mixed_external_import_aicore_invalid.pto create mode 100644 test/lit/vpto/backend_mixed_external_import_missing.pto create mode 100644 test/lit/vpto/backend_mixed_external_import_resolves.pto create mode 100644 test/lit/vpto/backend_mixed_external_import_signature_mismatch.pto create mode 100644 test/lit/vpto/backend_mixed_requires_output_file.pto create mode 100644 test/lit/vpto/backend_mixed_vpto_stub_signature_conflict.pto create mode 100644 test/lit/vpto/backend_public_abi_export_suffix.pto create mode 100644 test/lit/vpto/backend_reserved_aicore_abi_suffix_invalid.pto create mode 100644 test/lit/vpto/backend_reserved_public_abi_suffix_invalid.pto create mode 100644 test/lit/vpto/backend_vpto_child_missing_kernel_kind_invalid.pto create mode 100644 test/vpto/cases/micro-op/backend/mixed-external-vadd/compare.py create mode 100644 test/vpto/cases/micro-op/backend/mixed-external-vadd/golden.py create mode 100644 test/vpto/cases/micro-op/backend/mixed-external-vadd/kernel.pto create mode 100644 test/vpto/cases/micro-op/backend/mixed-external-vadd/launch.cpp create mode 100644 test/vpto/cases/micro-op/backend/mixed-external-vadd/main.cpp create mode 100644 test/vpto/cases/micro-op/backend/mixed-external-vadd/ptoas.flags rename tools/ptoas/{VPTOFatobjEmission.cpp => ObjectEmission.cpp} (58%) create mode 100644 tools/ptoas/ObjectEmission.h create mode 100644 tools/ptoas/PTOASDriver.cpp rename tools/ptoas/{VPTOFatobjEmission.h => PTOASDriver.h} (57%) diff --git a/docs/designs/module-backend-attr-and-mixed-fatobj.md b/docs/designs/module-backend-attr-and-mixed-fatobj.md new file mode 100644 index 0000000000..72abf4dd8b --- /dev/null +++ b/docs/designs/module-backend-attr-and-mixed-fatobj.md @@ -0,0 +1,553 @@ +# Module Backend Driver and Object Emission + +## Purpose + +This design adds a module-level backend selector and reorganizes `ptoas` around +an explicit driver layer. The driver owns command-line parsing, input loading, +PTO compilation planning, backend dispatch, Bisheng compilation, and final +fatobj linking. + +The design also replaces backend-specific object/fatobj emitters with one +`ObjectEmission` module. `ObjectEmission` provides both high-level helpers and +fine-grained stage APIs for compiling C++ or VPTO LLVM artifacts into device +objects and packing them into a fatobj. + +## User Contract + +### `pto.backend` Module Attribute + +Use the `pto.backend` attribute on `module`: + +```mlir +module attributes {pto.backend = "emitc"} { + ... +} + +module attributes {pto.backend = "vpto"} { + ... +} +``` + +Valid values are `emitc` and `vpto`. Unknown values are invalid. + +The attribute is intentionally module-level. Function-level backend selection is +not part of this contract; a function uses the backend of its nearest enclosing +backend module. + +### Backend Selection Priority + +`--pto-backend` remains the strongest selector. + +1. If the user passes `--pto-backend=emitc` or `--pto-backend=vpto`, the command + line forces that backend and `pto.backend` attributes are not used for backend + selection. +2. If the user does not pass `--pto-backend`, PTOAS reads `pto.backend` attributes. +3. If neither the command line nor the input specifies a backend, PTOAS keeps + the existing default: `emitc`. + +The driver must distinguish "the user did not pass `--pto-backend`" from "the +user passed `--pto-backend=emitc`"; the current option default alone is not +enough. + +### Single-Backend Input + +For a single module: + +```mlir +module attributes {pto.target_arch = "a5", pto.backend = "vpto"} { + func.func @kernel(...) attributes {pto.aicore} { + ... + } +} +``` + +If `--pto-backend` is absent, this is equivalent to: + +```bash +ptoas --pto-backend=vpto input.pto -o kernel.o +``` + +For `pto.backend = "emitc"`, PTOAS still uses the existing EmitC lowering to +produce CCE C++ source, but compilation of that source is internal to PTOAS. +The final output is a fatobj rather than externally compiled C++. + +### Mixed-Backend Container + +A mixed-backend input is an outer module containing backend-selected child +modules: + +```mlir +module attributes {pto.target_arch = "a5"} { + module attributes {pto.backend = "emitc"} { + func.func public @emitc_kernel(...) attributes {pto.aicore} { + ... + } + } + + module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind + } { + func.func public @mixed_kernel(...) { + ... + } + func.func @mixed_kernel_device(...) attributes {pto.aicore} { + ... + } + } + + module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind + } { + func.func public @mixed_kernel(...) { + ... + } + func.func @mixed_kernel_device(...) attributes {pto.aicore} { + ... + } + } +} +``` + +The outer module carries shared attributes such as `pto.target_arch`. Each child +module carries exactly one backend. Child modules may also carry backend-specific +attributes such as `pto.kernel_kind`. + +If child modules use more than one backend and `--pto-backend` is absent, PTOAS +enters mixed fatobj mode. + +If child modules use more than one backend and `--pto-backend` is present, the +command line wins. The input is treated as a forced single-backend compilation; +unsupported combinations should fail with a direct diagnostic rather than +silently compiling only part of the input. + +## Export and Symbol Contract + +For non-`pto.aicore` functions, PTOAS follows the normal `func.func` shape: + +| `func.func` form | Meaning in a backend child module | +|------------------|-----------------------------------| +| public function with a body | exported definition | +| private function with a body | module-local helper definition | +| private function without a body | external import declaration | + +A function without a body does not export a symbol from the current child +module. It declares a symbol that must be resolved by another backend child +module or by a later link input. A function with a body defines the symbol in +the current child module; it is exported only when it is public. + +For `pto.aicore` functions, symbol visibility is ignored by this export +contract. A `pto.aicore` function may be public or private in MLIR symbol terms, +but that visibility does not decide whether it is exported. + +Users write source-level function names without the `.vector` or `.cube` ABI +suffix. PTOAS derives the exported ABI name for every public non-`pto.aicore` +definition in a VPTO child module by appending the suffix that matches the +module's `pto.kernel_kind`: + +| `pto.kernel_kind` | Source symbol | Generated ABI export symbol | +|-------------------|---------------|-----------------------------| +| `#pto.kernel_kind` | `@foo` | `@foo.vector` | +| `#pto.kernel_kind` | `@foo` | `@foo.cube` | + +`pto.backend = "emitc"` modules follow the same source-level rule: users do not add +`.vector` or `.cube` in the PTO input. The CCE frontend or PTOAS object path +adds the required suffix for EmitC-generated CCE source. + +Users also write `pto.aicore` function names without the `_mix_aiv` or +`_mix_aic` device ABI suffix. PTOAS derives the lowered device ABI name from +the source symbol and the target unit: + +| `pto.kernel_kind` | Source `pto.aicore` symbol | Generated device ABI symbol | +|-------------------|----------------------------|-----------------------------| +| `#pto.kernel_kind` | `@foo_device` | `@foo_device_mix_aiv` | +| `#pto.kernel_kind` | `@foo_device` | `@foo_device_mix_aic` | + +For VPTO modules, PTOAS therefore tracks two names for an exported entry: + +1. the public non-`pto.aicore` source function name written in the input, such + as `@foo` +2. the source `pto.aicore` device function name written in the input, such as + `@foo_device` +3. the generated public ABI export name, such as `@foo.vector` or `@foo.cube` +4. the generated lowered device ABI name, such as `@foo_device_mix_aiv` or + `@foo_device_mix_aic` + +Any aliasing or export metadata needed to connect the public `.vector` / +`.cube` export name to the lowered device symbol is an implementation detail of +the fatobj path. + +### Mixed-Backend External Call Case + +Cross-backend calls are represented as normal external symbol references inside +the caller module. The callee must be provided by exactly one exported +non-`pto.aicore` function in another backend child module. + +```mlir +module attributes {pto.target_arch = "a5"} { + module attributes {pto.backend = "emitc"} { + func.func private @vpto_post( + %src: !pto.ptr, + %dst: !pto.ptr, + %n: index) + + func.func public @emitc_entry( + %src: !pto.ptr, + %dst: !pto.ptr, + %n: index) attributes {pto.aicore} { + func.call @vpto_post(%src, %dst, %n) + : (!pto.ptr, !pto.ptr, index) -> () + return + } + } + + module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind + } { + func.func public @vpto_post( + %src: !pto.ptr, + %dst: !pto.ptr, + %n: index) { + ... + return + } + } +} +``` + +In this case, the EmitC child module owns an external import declaration and the +call site for source symbol `@vpto_post`. The VPTO child module owns the public +exported definition for source symbol `@vpto_post`. During plan verification, +the driver resolves this as a cross-backend external call, records that the +callee's generated ABI symbol is `@vpto_post.vector`, and rewrites or forwards +that ABI name to the backend object path. The backend pipelines do not inline +or lower across the child-module boundary. + +Fatobj generation must allow unresolved device externals by passing `-dc` to +the fatobj repack step. The reference is resolved when the generated fatobj is +linked with the final host/kernel binary; that downstream Bisheng link must pass +`--cce-fatobj-link`. + +The declaration is private because the current upstream `func.func` verifier +rejects body-less public declarations. + +## Validation + +PTOAS should reject: + +1. Any `pto.backend` value other than `emitc` or `vpto`. +2. A mixed-backend container whose child module is missing `pto.backend`. +3. A `vpto` child module missing `pto.kernel_kind`. +4. A source-level public non-`pto.aicore` function in a VPTO child module that + already uses the reserved `.vector` or `.cube` ABI suffix. +5. A generated public ABI export name whose suffix does not match the VPTO + child module's `pto.kernel_kind`. +6. A source-level `pto.aicore` function in a VPTO child module that already + uses the reserved `_mix_aiv` or `_mix_aic` device ABI suffix. +7. A generated lowered device ABI name whose suffix does not match the VPTO + child module's `pto.kernel_kind`. +8. A cross-backend external call whose callee has no matching public + non-`pto.aicore` definition in another backend child module. +9. A cross-backend external call whose callee matches multiple exports or whose + signature does not match the exported definition. +10. A body-less external declaration that is not private or is marked + `pto.aicore`. + +## Architecture + +```text +ptoas main + └─ PTOASDriver::run(argc, argv) + ├─ parse command line + ├─ load input + ├─ setup MLIR context + ├─ parse .pto or decode .ptobc + ├─ buildPlan() + ├─ verifyPlan() + ├─ compilePTO() + │ ├─ emitc child ── PTO pipeline ── CCE source + │ ├─ vpto vector ─ PTO pipeline ── VPTO LLVM + │ └─ vpto cube ─── PTO pipeline ── VPTO LLVM + ├─ compileDeviceObjects() ── calls ObjectEmission + │ ├─ cpp source -> dav-c310-vec / dav-c310-cube objects + │ └─ VPTO LLVM -> dav-c310-vec / dav-c310-cube objects + ├─ linkFatobj() ── calls ObjectEmission fatobj stages + └─ write -o +``` + +Pass pipelines stop at compiler artifacts: + +| Backend path | Pipeline output | +|--------------|-----------------| +| EmitC | CCE C++ source | +| VPTO | LLVM module | + +`ObjectEmission` is not a pass and is not appended to any pass pipeline. It is a +driver-called service layer. The driver chooses which `ObjectEmission` stage +APIs to call after the relevant pipeline has returned its artifact. + +## `PTOASDriver` + +`ptoas` enters the driver layer immediately. Command-line parsing, input +loading, MLIR context setup, textual `.pto` parsing, and `.ptobc` decoding are +all driver responsibilities. Single-backend EmitC, single-backend VPTO, and +mixed-backend fatobj all go through this driver layer; backend-specific +shortcuts should not bypass it. + +The driver should be implemented as a separate component, for example +`PTOASDriver.{h,cpp}`, rather than growing `tools/ptoas/ptoas.cpp` further. +`ptoas.cpp` should become a thin process entrypoint that registers the version +printer if needed and calls `PTOASDriver::run(argc, argv)`. + +### Driver Responsibilities + +The driver owns the full `run()` flow: + +```text +setupAndPlan() + -> verifyPlan() + -> compilePTO() + -> compileDeviceObjects() + -> linkFatobj() + -> writeOutput() +``` + +The three compile/link stages are the core artifact-production phases, but they +do not replace setup and planning. Command-line parsing, input loading, MLIR +setup, parse/decode, backend resolution, module planning, and plan verification +remain driver work before artifact production starts. + +1. `setupAndPlan()`. + - Parse PTOAS user-facing options from `argc` / `argv`. + - Track whether `--pto-backend` appeared on the command line while parsing. + - Load `.pto` text or `.ptobc` bytes. + - Register and load the dialects needed by PTOAS. + - Decode `.ptobc` inputs. + - Parse textual `.pto` inputs with the effective parser target arch. + - Read `pto.backend` attributes when the CLI does not force a backend. + - Decide single-backend EmitC, single-backend VPTO, or mixed-backend fatobj + mode. + - Preserve shared outer attributes such as `pto.target_arch`. + - Identify backend child modules. + - Clone or isolate child modules before backend-specific passes mutate them. + +2. `verifyPlan()`. + - Validate `pto.backend` values and child module shape. + - Validate VPTO `pto.kernel_kind` requirements. + - Reject source-level public non-`pto.aicore` names that already use + reserved `.vector` or `.cube` ABI suffixes. + - Reject source-level `pto.aicore` names that already use reserved + `_mix_aiv` or `_mix_aic` device ABI suffixes. + - Derive and validate generated public ABI export names for VPTO children. + - Derive and validate generated lowered device ABI names for VPTO children. + - Collect source symbols and generated ABI export symbols from all backend + child modules. + - Resolve external calls across backend child modules by source symbol name + and signature, then record the generated ABI callee symbol. + +3. `compilePTO()`. + - Shared PTO planning/lowering remains in the driver-controlled flow. + - EmitC children are compiled to CCE C++ source. + - VPTO children are compiled to VPTO LLVM modules. + - The pass pipelines do not invoke Bisheng, object emission, or fatobj + linking as a tail action. + +4. `compileDeviceObjects()`. + - Call `ObjectEmission` for CCE C++ source -> device objects. + - Call `ObjectEmission` for VPTO LLVM modules -> device objects. + - Collect `BackendArtifact` results. + +5. `linkFatobj()`. + - Call `ObjectEmission` for device object merge, host stub compilation, and + final fatobj repack. + - Repack the fatobj with `-dc` so cross-fatobj device externals can remain + relocatable until the final `--cce-fatobj-link` link. + +6. `writeOutput()`. + - In fatobj mode, write the final packed object to `-o`. + - In debug/IR-print modes, route output according to the existing debug flag + semantics. + +Current implementation status: + +- The driver parses `pto.backend`, builds an explicit plan, clones backend child + modules, propagates shared outer attributes such as `pto.target_arch`, and + dispatches mixed children into backend-specific compiler-artifact generation. +- `ObjectEmission` owns the CCE/Bisheng object and fatobj stage APIs. +- Mixed fatobj mode requires an explicit `-o` file path and rejects debug IR + output flags. +- Cross-backend external imports are represented by private body-less + `func.func` declarations; the driver resolves them to exactly one exported + definition in another backend child. +- Host stub emission for multiple VPTO children is generated from the compiled + child modules as one merged stub source; duplicate logical kernel signatures + must agree. + +### Driver Data Model + +The driver should build an explicit compilation plan before running backend +work: + +```text +DriverConfig + argc / argv + input filename + output filename + target arch + build level + explicit CLI backend, if any + debug output flags + toolchain configuration + +CompilationPlan + mode: emitc | vpto | mixed-fatobj + modules: + - backend: emitc | vpto + kernel kind: none | vector | cube + module op clone + source exports + ABI exports + source device symbols + device ABI symbols + external references +``` + +Each completed backend job returns device artifacts: + +```text +BackendArtifact + backend: emitc | vpto + kind: vec-object | cube-object | generic-object + path: temporary object path + exported ABI symbols +``` + +The `ObjectEmission` fatobj stages consume `BackendArtifact` object paths. They +should not inspect MLIR modules. + +## `ObjectEmission` + +`ObjectEmission` owns all Bisheng-facing object and fatobj operations. It +replaces the separate C++ and VPTO fatobj emitter concepts with one module that +supports both high-level emit calls and fine-grained stage calls. + +`ObjectEmission` does not decide backend selection and does not run PTO or VPTO +MLIR lowering pipelines. The driver produces C++ source and VPTO LLVM modules, +then requests object/fatobj operations from this component. + +### High-Level Device Emit Interfaces + +```text +emitCppVectorDeviceObject(cppSource) -> dav-c310-vec object +emitCppCubeDeviceObject(cppSource) -> dav-c310-cube object +emitVPTOVectorDeviceObject(llvmModule) -> dav-c310-vec object +emitVPTOCubeDeviceObject(llvmModule) -> dav-c310-cube object +``` + +The C++ source path compiles the same EmitC-generated source for both device +targets when requested: + +```text +CCE C++ source + ├─ Bisheng -xcce --cce-aicore-arch=dav-c310-vec ── vector device object + └─ Bisheng -xcce --cce-aicore-arch=dav-c310-cube ── cube device object +``` + +The VPTO path compiles VPTO LLVM modules for the matching device target: + +```text +VPTO vector LLVM ── Bisheng -x ir --cce-aicore-arch=dav-c310-vec ── vector object +VPTO cube LLVM ─── Bisheng -x ir --cce-aicore-arch=dav-c310-cube ── cube object +``` + +### Fine-Grained Stage API + +`ObjectEmission` should expose stage-level APIs so the driver, tests, and debug +tools can run individual pieces without going through a monolithic fatobj call: + +```text +writeCppSource(cppSource) -> path +writeLLVMModule(llvmModule) -> path + +compileCppToDeviceObject(cppPath, target: vec|cube) -> object path +compileLLVMToDeviceObject(llPath, target: vec|cube) -> object path + +mergeDeviceObjects(objectPaths) -> merged object path +writeHostStubSource(stubSource) -> path +compileHostStubToObject(stubPath, mergedObjectPath, moduleId) -> host object path +repackFatobj(hostObjectPath, moduleId, outputPath) -> output fatobj +``` + +The high-level emit helpers are thin compositions of these stage APIs. For +example, `emitCppVectorDeviceObject` is `writeCppSource` followed by +`compileCppToDeviceObject(..., vec)`. A full fatobj build is +`mergeDeviceObjects` followed by host stub compilation and repack. + +### ObjectEmission Responsibilities + +1. Discover and validate the Bisheng/cce-ld/ld.lld toolchain. +2. Own temporary-file creation and cleanup for source, LLVM IR, device objects, + command stderr, merged object, host stub source, and host stub object. +3. Compile CCE C++ source to vector and/or cube device objects. +4. Compile VPTO LLVM modules to vector and/or cube device objects. +5. Merge an arbitrary list of device objects. +6. Compile the host stub with the merged device object embedded. +7. Repack the final fatobj to `-o` with `-dc`; the final consumer that links + this fatobj into a host/kernel binary must use `--cce-fatobj-link`. +8. Keep diagnostics separated by stage and artifact kind. + +## Implementation Plan + +1. Introduce `PTOASDriver`. + - Move command-line parsing, input loading, MLIR context setup, parse/decode, + backend selection, module planning, backend dispatch, Bisheng compile/link + orchestration, and final fatobj output decisions behind the driver. + - Keep `tools/ptoas/ptoas.cpp` as the thin entrypoint that calls + `PTOASDriver::run(argc, argv)`. + +2. Add `pto.backend` attribute handling. + - Parse and validate `pto.backend`. + - Track whether `--pto-backend` was explicitly provided. + - Resolve single-backend or mixed-backend mode in the driver. + +3. Add mixed-container planning. + - Identify backend child modules. + - Clone/isolate child modules before mutation. + - Preserve shared outer attributes such as `pto.target_arch`. + - Record source-level exported non-`pto.aicore` symbols, generated ABI export + symbols, source-level `pto.aicore` symbols, and generated device ABI + symbols. + +4. Refactor backend pipeline entrypoints. + - EmitC pipeline returns CCE C++ source. + - VPTO pipeline returns LLVM module(s). + - Object/fatobj emission is not appended as pipeline finalization. + +5. Introduce `ObjectEmission`. + - Move generic toolchain discovery, temporary-file management, object + compilation, object merge, host stub compilation, and repack operations + into this module. + - Provide cpp/vpto vector/cube high-level emit helpers. + - Provide fine-grained stage APIs. + +6. Route final output through the driver. + - Driver collects `BackendArtifact` results. + - Driver calls `ObjectEmission` merge/host-stub/repack stages. + - Driver writes or keeps the final `-o` result. + +7. Add focused tests. + - `pto.backend` attr fallback when `--pto-backend` is absent. + - CLI backend override when `--pto-backend` is present. + - Rejection of source-level public non-`pto.aicore` symbols that already use + reserved `.vector` / `.cube` ABI suffixes. + - Generation of VPTO public non-`pto.aicore` `.vector` / `.cube` ABI export + symbols from suffix-free source symbols. + - Rejection of source-level `pto.aicore` symbols that already use reserved + `_mix_aiv` / `_mix_aic` device ABI suffixes. + - Generation of lowered `pto.aicore` `_mix_aiv` / `_mix_aic` device ABI + symbols from suffix-free source symbols. + - Driver-planned mixed backend mode. + - ObjectEmission stage APIs through tests that do not require full pipeline + execution when possible. diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index c664712446..cec6676f18 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -2740,14 +2740,18 @@ struct FuncToEmitC : public OpConversionPattern { } if (op.isDeclaration()) { - emitcFunc.setSpecifiersAttr(rewriter.getStrArrayAttr({"extern"})); + if (op->hasAttr("pto.external_abi")) + emitcFunc.setSpecifiersAttr( + rewriter.getStrArrayAttr({"extern \"C\"", "AICORE"})); + else + emitcFunc.setSpecifiersAttr(rewriter.getStrArrayAttr({"extern"})); rewriter.eraseOp(op); return success(); } if (pto::isPTOEntryFunction(op)) { emitcFunc.setSpecifiersAttr( - rewriter.getStrArrayAttr({"__global__ AICORE"})); + rewriter.getStrArrayAttr({"extern \"C\"", "__global__ AICORE"})); } else if (op.isPrivate()) { emitcFunc.setSpecifiersAttr( rewriter.getStrArrayAttr({"static", "AICORE"})); diff --git a/test/lit/vpto/backend_attr_invalid.pto b/test/lit/vpto/backend_attr_invalid.pto new file mode 100644 index 0000000000..0aa9004f00 --- /dev/null +++ b/test/lit/vpto/backend_attr_invalid.pto @@ -0,0 +1,21 @@ +// 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: not ptoas --pto-arch=a5 --emit-vpto %s -o - 2>&1 | FileCheck %s + +module attributes { + pto.backend = "other", + pto.kernel_kind = #pto.kernel_kind, + pto.target_arch = "a5" +} { + func.func @backend_attr_invalid() { + return + } +} + +// CHECK: invalid pto.backend 'other'. Expected 'emitc' or 'vpto'. diff --git a/test/lit/vpto/backend_attr_selects_vpto.pto b/test/lit/vpto/backend_attr_selects_vpto.pto new file mode 100644 index 0000000000..1d6997f5b7 --- /dev/null +++ b/test/lit/vpto/backend_attr_selects_vpto.pto @@ -0,0 +1,26 @@ +// 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 --emit-vpto %s -o - | FileCheck %s + +module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind, + pto.target_arch = "a5" +} { + func.func @backend_attr_selects_vpto() { + %ctrl = pto.get_ctrl : i64 + pto.set_ctrl %ctrl : i64 + return + } +} + +// CHECK: module +// CHECK-SAME: pto.backend = "vpto" +// CHECK: pto.get_ctrl +// CHECK: pto.set_ctrl diff --git a/test/lit/vpto/backend_child_attr_inherits_outer_arch.pto b/test/lit/vpto/backend_child_attr_inherits_outer_arch.pto new file mode 100644 index 0000000000..d92a13afbc --- /dev/null +++ b/test/lit/vpto/backend_child_attr_inherits_outer_arch.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 --emit-vpto %s -o - | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind + } { + func.func @backend_child_attr_inherits_outer_arch() { + %ctrl = pto.get_ctrl : i64 + pto.set_ctrl %ctrl : i64 + return + } + } +} + +// CHECK: module +// CHECK-SAME: pto.target_arch = "a5" +// CHECK: module +// CHECK-SAME: pto.backend = "vpto" +// CHECK: pto.get_ctrl +// CHECK: pto.set_ctrl diff --git a/test/lit/vpto/backend_child_attr_selects_vpto.pto b/test/lit/vpto/backend_child_attr_selects_vpto.pto new file mode 100644 index 0000000000..8f817c8bab --- /dev/null +++ b/test/lit/vpto/backend_child_attr_selects_vpto.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 --emit-vpto %s -o - | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind + } { + func.func @backend_child_attr_selects_vpto() { + %ctrl = pto.get_ctrl : i64 + pto.set_ctrl %ctrl : i64 + return + } + } +} + +// CHECK: module +// CHECK-SAME: pto.target_arch = "a5" +// CHECK: module +// CHECK-SAME: pto.backend = "vpto" +// CHECK: pto.get_ctrl +// CHECK: pto.set_ctrl diff --git a/test/lit/vpto/backend_cli_override.pto b/test/lit/vpto/backend_cli_override.pto new file mode 100644 index 0000000000..45ebf88ec9 --- /dev/null +++ b/test/lit/vpto/backend_cli_override.pto @@ -0,0 +1,23 @@ +// 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: not ptoas --pto-arch=a5 --pto-backend=emitc --emit-vpto %s -o - 2>&1 | FileCheck %s + +module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind, + pto.target_arch = "a5" +} { + func.func @backend_cli_override() { + %ctrl = pto.get_ctrl : i64 + pto.set_ctrl %ctrl : i64 + return + } +} + +// CHECK: Error: VPTO-specific flags require --pto-backend=vpto or pto.backend = "vpto". diff --git a/test/lit/vpto/backend_cli_override_to_vpto.pto b/test/lit/vpto/backend_cli_override_to_vpto.pto new file mode 100644 index 0000000000..f3031ba631 --- /dev/null +++ b/test/lit/vpto/backend_cli_override_to_vpto.pto @@ -0,0 +1,24 @@ +// 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-backend=vpto --emit-vpto %s -o - | FileCheck %s + +module attributes { + pto.backend = "emitc", + pto.kernel_kind = #pto.kernel_kind, + pto.target_arch = "a5" +} { + func.func @backend_cli_override_to_vpto() { + %ctrl = pto.get_ctrl : i64 + pto.set_ctrl %ctrl : i64 + return + } +} + +// CHECK: pto.get_ctrl +// CHECK: pto.set_ctrl diff --git a/test/lit/vpto/backend_mixed_child_attrs_rejected.pto b/test/lit/vpto/backend_mixed_child_attrs_rejected.pto new file mode 100644 index 0000000000..8367461871 --- /dev/null +++ b/test/lit/vpto/backend_mixed_child_attrs_rejected.pto @@ -0,0 +1,28 @@ +// 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: not ptoas --pto-arch=a5 --emit-vpto %s -o - 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.backend = "emitc"} { + func.func @emitc_child() attributes {pto.aicore} { + return + } + } + + module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind + } { + func.func @vpto_child() { + return + } + } +} + +// CHECK: mixed pto.backend fatobj mode does not support debug IR output flags diff --git a/test/lit/vpto/backend_mixed_child_missing_backend_invalid.pto b/test/lit/vpto/backend_mixed_child_missing_backend_invalid.pto new file mode 100644 index 0000000000..d54e910149 --- /dev/null +++ b/test/lit/vpto/backend_mixed_child_missing_backend_invalid.pto @@ -0,0 +1,34 @@ +// 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: not ptoas --pto-arch=a5 --emit-vpto %s -o - 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.backend = "emitc"} { + func.func @emitc_child() attributes {pto.aicore} { + return + } + } + + module { + func.func @missing_backend_child() { + return + } + } + + module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind + } { + func.func @vpto_child() { + return + } + } +} + +// CHECK: mixed-backend container child module is missing pto.backend diff --git a/test/lit/vpto/backend_mixed_duplicate_abi_export_invalid.pto b/test/lit/vpto/backend_mixed_duplicate_abi_export_invalid.pto new file mode 100644 index 0000000000..d1bfdf21c5 --- /dev/null +++ b/test/lit/vpto/backend_mixed_duplicate_abi_export_invalid.pto @@ -0,0 +1,37 @@ +// 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: not ptoas --pto-arch=a5 --emit-vpto %s -o - 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.backend = "emitc"} { + func.func @emitc_child() attributes {pto.aicore} { + return + } + } + + module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind + } { + func.func public @vpto_post() { + return + } + } + + module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind + } { + func.func public @vpto_post() { + return + } + } +} + +// CHECK: generated ABI export symbol 'vpto_post.vector' is defined by multiple backend child modules diff --git a/test/lit/vpto/backend_mixed_external_import_aicore_invalid.pto b/test/lit/vpto/backend_mixed_external_import_aicore_invalid.pto new file mode 100644 index 0000000000..77fa81d662 --- /dev/null +++ b/test/lit/vpto/backend_mixed_external_import_aicore_invalid.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: not ptoas --pto-arch=a5 %s -o %t.o 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.backend = "emitc"} { + func.func private @vpto_post() attributes {pto.aicore} + } + + module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind + } { + func.func public @vpto_post() { + return + } + func.func @vpto_device() attributes {pto.aicore} { + return + } + } +} + +// CHECK: external func.func declaration must not be marked pto.aicore diff --git a/test/lit/vpto/backend_mixed_external_import_missing.pto b/test/lit/vpto/backend_mixed_external_import_missing.pto new file mode 100644 index 0000000000..d5e2f59f12 --- /dev/null +++ b/test/lit/vpto/backend_mixed_external_import_missing.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: not ptoas --pto-arch=a5 %s -o %t.o 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.backend = "emitc"} { + func.func private @vpto_post(%arg0: i64) + + func.func @emitc_entry() attributes {pto.aicore} { + %c0 = arith.constant 0 : i64 + func.call @vpto_post(%c0) : (i64) -> () + return + } + } + + module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind + } { + func.func @vpto_device() attributes {pto.aicore} { + return + } + } +} + +// CHECK: cross-backend external import 'vpto_post' has no matching public non-pto.aicore definition in another backend child module diff --git a/test/lit/vpto/backend_mixed_external_import_resolves.pto b/test/lit/vpto/backend_mixed_external_import_resolves.pto new file mode 100644 index 0000000000..5357a2b393 --- /dev/null +++ b/test/lit/vpto/backend_mixed_external_import_resolves.pto @@ -0,0 +1,36 @@ +// 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: not env -u ASCEND_HOME_PATH ptoas --pto-arch=a5 %s -o %t.o 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.backend = "emitc"} { + func.func private @vpto_post(%arg0: i64) + + func.func @emitc_entry() attributes {pto.aicore} { + %c0 = arith.constant 0 : i64 + func.call @vpto_post(%c0) : (i64) -> () + return + } + } + + module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind + } { + func.func public @vpto_post(%arg0: i64) { + return + } + func.func @vpto_device() attributes {pto.aicore} { + return + } + } +} + +// CHECK-NOT: cross-backend external import +// CHECK: ASCEND_HOME_PATH is required diff --git a/test/lit/vpto/backend_mixed_external_import_signature_mismatch.pto b/test/lit/vpto/backend_mixed_external_import_signature_mismatch.pto new file mode 100644 index 0000000000..71d55f0145 --- /dev/null +++ b/test/lit/vpto/backend_mixed_external_import_signature_mismatch.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: not ptoas --pto-arch=a5 %s -o %t.o 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.backend = "emitc"} { + func.func private @vpto_post(%arg0: i64) + } + + module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind + } { + func.func public @vpto_post(%arg0: i32) { + return + } + func.func @vpto_device() attributes {pto.aicore} { + return + } + } +} + +// CHECK: cross-backend external import 'vpto_post' signature does not match exported definition diff --git a/test/lit/vpto/backend_mixed_requires_output_file.pto b/test/lit/vpto/backend_mixed_requires_output_file.pto new file mode 100644 index 0000000000..3350f8dbee --- /dev/null +++ b/test/lit/vpto/backend_mixed_requires_output_file.pto @@ -0,0 +1,28 @@ +// 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: not ptoas --pto-arch=a5 %s -o - 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.backend = "emitc"} { + func.func @emitc_child() attributes {pto.aicore} { + return + } + } + + module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind + } { + func.func @vpto_child() { + return + } + } +} + +// CHECK: mixed pto.backend fatobj mode requires an explicit file path passed with -o diff --git a/test/lit/vpto/backend_mixed_vpto_stub_signature_conflict.pto b/test/lit/vpto/backend_mixed_vpto_stub_signature_conflict.pto new file mode 100644 index 0000000000..dfdd4a423b --- /dev/null +++ b/test/lit/vpto/backend_mixed_vpto_stub_signature_conflict.pto @@ -0,0 +1,37 @@ +// 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: not ptoas --pto-arch=a5 %s -o %t.o 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.backend = "emitc"} { + func.func @emitc_child() attributes {pto.aicore} { + return + } + } + + module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind + } { + func.func @same_kernel(%arg0: i32) attributes {pto.aicore} { + return + } + } + + module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind + } { + func.func @same_kernel(%arg0: i64) attributes {pto.aicore} { + return + } + } +} + +// CHECK: mixed kernel variants disagree on host stub signature for 'same_kernel' diff --git a/test/lit/vpto/backend_public_abi_export_suffix.pto b/test/lit/vpto/backend_public_abi_export_suffix.pto new file mode 100644 index 0000000000..09e8e8d7de --- /dev/null +++ b/test/lit/vpto/backend_public_abi_export_suffix.pto @@ -0,0 +1,38 @@ +// 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 --emit-vpto %s -o - | FileCheck %s + +module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind, + pto.target_arch = "a5" +} { + func.func public @vpto_post(%value: i64) { + func.call @local_helper(%value) : (i64) -> () + return + } + + func.func private @local_helper(%value: i64) { + return + } + + func.func @device_kernel() attributes {pto.aicore} { + %c0 = arith.constant 0 : i64 + func.call @vpto_post(%c0) : (i64) -> () + return + } +} + +// CHECK: pto.kernel_kind = #pto.kernel_kind +// CHECK: func.func public @vpto_post.vector +// CHECK: call @local_helper +// CHECK: func.func private @local_helper +// CHECK: func.func @device_kernel +// CHECK: call @vpto_post.vector +// CHECK-NOT: call @vpto_post( diff --git a/test/lit/vpto/backend_reserved_aicore_abi_suffix_invalid.pto b/test/lit/vpto/backend_reserved_aicore_abi_suffix_invalid.pto new file mode 100644 index 0000000000..c57562a709 --- /dev/null +++ b/test/lit/vpto/backend_reserved_aicore_abi_suffix_invalid.pto @@ -0,0 +1,21 @@ +// 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: not ptoas --pto-arch=a5 --emit-vpto %s -o - 2>&1 | FileCheck %s + +module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind, + pto.target_arch = "a5" +} { + func.func @device_kernel_mix_aic() attributes {pto.aicore} { + return + } +} + +// CHECK: source-level pto.aicore function 'device_kernel_mix_aic' must not use reserved device ABI suffix '_mix_aic' diff --git a/test/lit/vpto/backend_reserved_public_abi_suffix_invalid.pto b/test/lit/vpto/backend_reserved_public_abi_suffix_invalid.pto new file mode 100644 index 0000000000..13a03fe5b0 --- /dev/null +++ b/test/lit/vpto/backend_reserved_public_abi_suffix_invalid.pto @@ -0,0 +1,21 @@ +// 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: not ptoas --pto-arch=a5 --emit-vpto %s -o - 2>&1 | FileCheck %s + +module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind, + pto.target_arch = "a5" +} { + func.func public @vpto_post.vector() { + return + } +} + +// CHECK: source-level public non-pto.aicore function 'vpto_post.vector' must not use reserved VPTO ABI suffix '.vector' diff --git a/test/lit/vpto/backend_vpto_child_missing_kernel_kind_invalid.pto b/test/lit/vpto/backend_vpto_child_missing_kernel_kind_invalid.pto new file mode 100644 index 0000000000..425925de04 --- /dev/null +++ b/test/lit/vpto/backend_vpto_child_missing_kernel_kind_invalid.pto @@ -0,0 +1,19 @@ +// 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: not ptoas --pto-arch=a5 --emit-vpto %s -o - 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.backend = "vpto"} { + func.func @missing_kernel_kind() { + return + } + } +} + +// CHECK: vpto backend child module must carry 'pto.kernel_kind' diff --git a/test/vpto/cases/micro-op/backend/mixed-external-vadd/compare.py b/test/vpto/cases/micro-op/backend/mixed-external-vadd/compare.py new file mode 100644 index 0000000000..792d7644d5 --- /dev/null +++ b/test/vpto/cases/micro-op/backend/mixed-external-vadd/compare.py @@ -0,0 +1,40 @@ +#!/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. + +import os +import sys + +import numpy as np + + +def compare_bin(golden_path, output_path, dtype, eps): + if not os.path.exists(golden_path) or not os.path.exists(output_path): + return False + golden = np.fromfile(golden_path, dtype=dtype) + output = np.fromfile(output_path, dtype=dtype) + return ( + golden.shape == output.shape + and np.allclose(golden, output, atol=eps, rtol=eps, equal_nan=True) + ) + + +def main(): + strict = os.getenv("COMPARE_STRICT", "1") != "0" + ok = compare_bin("golden_v3.bin", "v3.bin", np.float32, 1e-4) + if not ok: + if strict: + print("[ERROR] compare failed") + sys.exit(2) + print("[WARN] compare failed (non-gating)") + return + print("[INFO] compare passed") + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/micro-op/backend/mixed-external-vadd/golden.py b/test/vpto/cases/micro-op/backend/mixed-external-vadd/golden.py new file mode 100644 index 0000000000..af261b5830 --- /dev/null +++ b/test/vpto/cases/micro-op/backend/mixed-external-vadd/golden.py @@ -0,0 +1,43 @@ +#!/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. + +import argparse +from pathlib import Path + +import numpy as np + + +ELEM_COUNT = 1024 +SEED = 29 + + +def generate(output_dir: Path, seed: int) -> None: + rng = np.random.default_rng(seed) + lhs = rng.uniform(-8.0, 8.0, size=(ELEM_COUNT,)).astype(np.float32) + rhs = rng.uniform(-8.0, 8.0, size=(ELEM_COUNT,)).astype(np.float32) + out = np.zeros((ELEM_COUNT,), dtype=np.float32) + golden = (lhs + rhs).astype(np.float32, copy=False) + + output_dir.mkdir(parents=True, exist_ok=True) + lhs.tofile(output_dir / "v1.bin") + rhs.tofile(output_dir / "v2.bin") + out.tofile(output_dir / "v3.bin") + golden.tofile(output_dir / "golden_v3.bin") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir", type=Path, default=Path(".")) + parser.add_argument("--seed", type=int, default=SEED) + args = parser.parse_args() + generate(args.output_dir, args.seed) + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/micro-op/backend/mixed-external-vadd/kernel.pto b/test/vpto/cases/micro-op/backend/mixed-external-vadd/kernel.pto new file mode 100644 index 0000000000..197ace46c1 --- /dev/null +++ b/test/vpto/cases/micro-op/backend/mixed-external-vadd/kernel.pto @@ -0,0 +1,93 @@ +// 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. + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.backend = "emitc"} { + func.func private @external_vadd(%lhs: !pto.ptr, + %rhs: !pto.ptr, + %out: !pto.ptr) + + func.func @mixed_external_vadd_kernel(%lhs: !pto.ptr, + %rhs: !pto.ptr, + %out: !pto.ptr) + attributes {pto.aicore} { + pto.section.vector { + func.call @external_vadd(%lhs, %rhs, %out) + : (!pto.ptr, !pto.ptr, !pto.ptr) -> () + } + return + } + } + + module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind + } { + func.func public @external_vadd(%lhs_gm: !pto.ptr, + %rhs_gm: !pto.ptr, + %out_gm: !pto.ptr) { + %c0 = arith.constant 0 : index + %c64 = arith.constant 64 : index + %c1024 = arith.constant 1024 : index + %c0_i64 = arith.constant 0 : i64 + %c32_i64 = arith.constant 32 : i64 + %c128_i64 = arith.constant 128 : i64 + %c4096_i64 = arith.constant 4096 : i64 + %c8192_i64 = arith.constant 8192 : i64 + %c1024_i32 = arith.constant 1024 : i32 + + %ub_lhs = pto.castptr %c0_i64 : i64 -> !pto.ptr + %ub_rhs = pto.castptr %c4096_i64 : i64 -> !pto.ptr + %ub_out = pto.castptr %c8192_i64 : i64 -> !pto.ptr + + pto.get_buf "PIPE_MTE2", 0, 0 + pto.mte_gm_ub %lhs_gm, %ub_lhs, %c0_i64, %c128_i64 + nburst(%c32_i64, %c128_i64, %c128_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.rls_buf "PIPE_MTE2", 0, 0 + + pto.get_buf "PIPE_MTE2", 1, 0 + pto.mte_gm_ub %rhs_gm, %ub_rhs, %c0_i64, %c128_i64 + nburst(%c32_i64, %c128_i64, %c128_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.rls_buf "PIPE_MTE2", 1, 0 + + pto.get_buf "PIPE_V", 0, 0 + pto.get_buf "PIPE_V", 1, 0 + pto.get_buf "PIPE_V", 2, 0 + pto.vecscope { + %_:1 = scf.for %offset = %c0 to %c1024 step %c64 + iter_args(%remaining = %c1024_i32) -> (i32) { + %mask, %next_remaining = pto.plt_b32 %remaining + : i32 -> !pto.mask, i32 + %lhs = pto.vlds %ub_lhs[%offset] + : !pto.ptr -> !pto.vreg<64xf32> + %rhs = pto.vlds %ub_rhs[%offset] + : !pto.ptr -> !pto.vreg<64xf32> + %sum = pto.vadd %lhs, %rhs, %mask + : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask + -> !pto.vreg<64xf32> + pto.vsts %sum, %ub_out[%offset], %mask + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + scf.yield %next_remaining : i32 + } + } + pto.rls_buf "PIPE_V", 0, 0 + pto.rls_buf "PIPE_V", 1, 0 + pto.rls_buf "PIPE_V", 2, 0 + + pto.get_buf "PIPE_MTE3", 2, 0 + pto.mte_ub_gm %ub_out, %out_gm, %c128_i64 + nburst(%c32_i64, %c128_i64, %c128_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.rls_buf "PIPE_MTE3", 2, 0 + pto.barrier #pto.pipe + return + } + } +} diff --git a/test/vpto/cases/micro-op/backend/mixed-external-vadd/launch.cpp b/test/vpto/cases/micro-op/backend/mixed-external-vadd/launch.cpp new file mode 100644 index 0000000000..97c1e01cda --- /dev/null +++ b/test/vpto/cases/micro-op/backend/mixed-external-vadd/launch.cpp @@ -0,0 +1,25 @@ +// 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. + +#ifndef __VEC_SCOPE__ +#define __VEC_SCOPE__ +#endif + +#include +#ifndef __CPU_SIM +#include "acl/acl.h" +#endif + +extern "C" __global__ [aicore] void mixed_external_vadd_kernel( + __gm__ float *lhs, __gm__ float *rhs, __gm__ float *out); + +void LaunchMixed_external_vadd_kernel(float *lhs, float *rhs, float *out, + void *stream) { + mixed_external_vadd_kernel<<<1, nullptr, stream>>>( + (__gm__ float *)lhs, (__gm__ float *)rhs, (__gm__ float *)out); +} diff --git a/test/vpto/cases/micro-op/backend/mixed-external-vadd/main.cpp b/test/vpto/cases/micro-op/backend/mixed-external-vadd/main.cpp new file mode 100644 index 0000000000..f270a081e6 --- /dev/null +++ b/test/vpto/cases/micro-op/backend/mixed-external-vadd/main.cpp @@ -0,0 +1,118 @@ +// 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. + +#include "acl/acl.h" +#include "test_common.h" + +#include +#include + +using namespace PtoTestCommon; + +#define ACL_CHECK(expr) \ + do { \ + const aclError _ret = (expr); \ + if (_ret != ACL_SUCCESS) { \ + std::fprintf(stderr, "[ERROR] %s failed: %d (%s:%d)\n", #expr, \ + (int)_ret, __FILE__, __LINE__); \ + const char *_recent = aclGetRecentErrMsg(); \ + if (_recent != nullptr && _recent[0] != '\0') \ + std::fprintf(stderr, "[ERROR] RecentErrMsg: %s\n", _recent); \ + rc = 1; \ + goto cleanup; \ + } \ + } while (0) + +#define FILE_CHECK(expr, path) \ + do { \ + if (!(expr)) { \ + std::fprintf(stderr, "[ERROR] file operation failed: %s (%s:%d)\n", \ + path, __FILE__, __LINE__); \ + rc = 1; \ + goto cleanup; \ + } \ + } while (0) + +void LaunchMixed_external_vadd_kernel(float *lhs, float *rhs, float *out, + void *stream); + +int main() { + constexpr size_t elemCount = 1024; + constexpr size_t bufSize = elemCount * sizeof(float); + + float *lhsHost = nullptr; + float *rhsHost = nullptr; + float *outHost = nullptr; + float *lhsDevice = nullptr; + float *rhsDevice = nullptr; + float *outDevice = nullptr; + + int rc = 0; + bool aclInited = false; + bool deviceSet = false; + int deviceId = 0; + aclrtStream stream = nullptr; + size_t inputSize = 0; + + ACL_CHECK(aclInit(nullptr)); + aclInited = true; + if (const char *envDevice = std::getenv("ACL_DEVICE_ID")) + deviceId = std::atoi(envDevice); + ACL_CHECK(aclrtSetDevice(deviceId)); + deviceSet = true; + ACL_CHECK(aclrtCreateStream(&stream)); + + ACL_CHECK(aclrtMallocHost((void **)&lhsHost, bufSize)); + ACL_CHECK(aclrtMallocHost((void **)&rhsHost, bufSize)); + ACL_CHECK(aclrtMallocHost((void **)&outHost, bufSize)); + ACL_CHECK(aclrtMalloc((void **)&lhsDevice, bufSize, ACL_MEM_MALLOC_HUGE_FIRST)); + ACL_CHECK(aclrtMalloc((void **)&rhsDevice, bufSize, ACL_MEM_MALLOC_HUGE_FIRST)); + ACL_CHECK(aclrtMalloc((void **)&outDevice, bufSize, ACL_MEM_MALLOC_HUGE_FIRST)); + + inputSize = bufSize; + FILE_CHECK(ReadFile("./v1.bin", inputSize, lhsHost, bufSize) && + inputSize == bufSize, + "./v1.bin"); + inputSize = bufSize; + FILE_CHECK(ReadFile("./v2.bin", inputSize, rhsHost, bufSize) && + inputSize == bufSize, + "./v2.bin"); + inputSize = bufSize; + FILE_CHECK(ReadFile("./v3.bin", inputSize, outHost, bufSize) && + inputSize == bufSize, + "./v3.bin"); + + ACL_CHECK(aclrtMemcpy(lhsDevice, bufSize, lhsHost, bufSize, + ACL_MEMCPY_HOST_TO_DEVICE)); + ACL_CHECK(aclrtMemcpy(rhsDevice, bufSize, rhsHost, bufSize, + ACL_MEMCPY_HOST_TO_DEVICE)); + ACL_CHECK(aclrtMemcpy(outDevice, bufSize, outHost, bufSize, + ACL_MEMCPY_HOST_TO_DEVICE)); + + LaunchMixed_external_vadd_kernel(lhsDevice, rhsDevice, outDevice, stream); + ACL_CHECK(aclrtSynchronizeStream(stream)); + + ACL_CHECK(aclrtMemcpy(outHost, bufSize, outDevice, bufSize, + ACL_MEMCPY_DEVICE_TO_HOST)); + FILE_CHECK(WriteFile("./v3.bin", outHost, bufSize), "./v3.bin"); + +cleanup: + aclrtFree(lhsDevice); + aclrtFree(rhsDevice); + aclrtFree(outDevice); + aclrtFreeHost(lhsHost); + aclrtFreeHost(rhsHost); + aclrtFreeHost(outHost); + if (stream != nullptr) + aclrtDestroyStream(stream); + if (deviceSet) + aclrtResetDevice(deviceId); + if (aclInited) + aclFinalize(); + return rc; +} diff --git a/test/vpto/cases/micro-op/backend/mixed-external-vadd/ptoas.flags b/test/vpto/cases/micro-op/backend/mixed-external-vadd/ptoas.flags new file mode 100644 index 0000000000..48e3c5c2e6 --- /dev/null +++ b/test/vpto/cases/micro-op/backend/mixed-external-vadd/ptoas.flags @@ -0,0 +1 @@ +--pto-arch a5 diff --git a/test/vpto/scripts/run_host_vpto_validation.sh b/test/vpto/scripts/run_host_vpto_validation.sh index 0edac38825..2b6bdeee21 100755 --- a/test/vpto/scripts/run_host_vpto_validation.sh +++ b/test/vpto/scripts/run_host_vpto_validation.sh @@ -241,6 +241,7 @@ build_one_impl() { local launch_obj="${out_dir}/launch.o" local kernel_fatobj="${out_dir}/kernel.fatobj.o" local kernel_so="${out_dir}/lib${case_token}_kernel.so" + local -a ptoas_args=() [[ -f "${case_dir}/main.cpp" ]] || die "missing main.cpp for ${case_name}" [[ -f "${case_dir}/launch.cpp" ]] || die "missing launch.cpp for ${case_name}" @@ -249,8 +250,14 @@ build_one_impl() { [[ -f "${case_dir}/kernel.pto" ]] || die "missing kernel.pto for ${case_name}" + if [[ -f "${case_dir}/ptoas.flags" ]]; then + read -r -a ptoas_args < "${case_dir}/ptoas.flags" + else + read -r -a ptoas_args <<< "${PTOAS_FLAGS}" + fi + log "[$case_name] step 1/4: emit kernel fatobj" - "${PTOAS_BIN}" ${PTOAS_FLAGS} \ + "${PTOAS_BIN}" "${ptoas_args[@]}" \ "${case_dir}/kernel.pto" -o "${kernel_fatobj}" log "[$case_name] step 2/4: build launch object" diff --git a/tools/ptoas/CMakeLists.txt b/tools/ptoas/CMakeLists.txt index f25734d66a..11090c8a1c 100644 --- a/tools/ptoas/CMakeLists.txt +++ b/tools/ptoas/CMakeLists.txt @@ -20,8 +20,9 @@ set(LLVM_LINK_COMPONENTS # 原因：LLVM 构建目录里已经有一个 ptoas 了，重名会导致 CMake 报错。 add_llvm_executable(pto-opt ptoas.cpp + PTOASDriver.cpp VPTOHostStubEmission.cpp - VPTOFatobjEmission.cpp + ObjectEmission.cpp TilelangDaemon.cpp ) # ========================================================= diff --git a/tools/ptoas/VPTOFatobjEmission.cpp b/tools/ptoas/ObjectEmission.cpp similarity index 58% rename from tools/ptoas/VPTOFatobjEmission.cpp rename to tools/ptoas/ObjectEmission.cpp index 0fb060a124..87a4bb5d4b 100644 --- a/tools/ptoas/VPTOFatobjEmission.cpp +++ b/tools/ptoas/ObjectEmission.cpp @@ -6,12 +6,13 @@ // 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. -#include "VPTOFatobjEmission.h" +#include "ObjectEmission.h" #include "PTO/Transforms/VPTOLLVMEmitter.h" #include "llvm/ADT/SmallString.h" #include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/STLExtras.h" #include "llvm/IR/Module.h" #include "llvm/Support/Error.h" #include "llvm/Support/FileSystem.h" @@ -120,6 +121,13 @@ static std::optional getAscendHomePath() { return std::string(env); } +static std::optional getEnvPath(llvm::StringRef name) { + const char *env = std::getenv(name.str().c_str()); + if (!env || !*env) + return std::nullopt; + return std::string(env); +} + static std::string joinPath(llvm::StringRef lhs, llvm::StringRef rhs) { llvm::SmallString<256> joined(lhs); llvm::sys::path::append(joined, rhs); @@ -135,6 +143,59 @@ static std::optional locateProgram(llvm::StringRef envPath, return std::nullopt; } +static bool hasPTOISAHeader(llvm::StringRef includeDir) { + return llvm::sys::fs::exists(joinPath(includeDir, "pto/pto-inst.hpp")); +} + +static void addExistingIncludeDir(llvm::SmallVectorImpl &dirs, + llvm::StringRef path) { + if (path.empty() || !llvm::sys::fs::is_directory(path)) + return; + if (llvm::is_contained(dirs, path)) + return; + dirs.push_back(path.str()); +} + +static void addPTOISAIncludeDirs(llvm::SmallVectorImpl &dirs, + llvm::StringRef ptoIsaPath) { + if (ptoIsaPath.empty() || !llvm::sys::fs::is_directory(ptoIsaPath)) + return; + std::string includeDir = joinPath(ptoIsaPath, "include"); + if (hasPTOISAHeader(includeDir)) + addExistingIncludeDir(dirs, includeDir); + std::string commonDir = joinPath(ptoIsaPath, "tests/common"); + addExistingIncludeDir(dirs, commonDir); + if (hasPTOISAHeader(ptoIsaPath)) + addExistingIncludeDir(dirs, ptoIsaPath); +} + +static llvm::SmallVector +discoverCppIncludeDirs(llvm::StringRef ascendHome, + llvm::raw_ostream &diagOS, + std::string &ptoIsaPath) { + llvm::SmallVector includeDirs; + if (auto env = getEnvPath("PTO_ISA_PATH")) + ptoIsaPath = *env; + else if (auto env = getEnvPath("PTO_ISA_ROOT")) + ptoIsaPath = *env; + + addPTOISAIncludeDirs(includeDirs, ptoIsaPath); + addExistingIncludeDir(includeDirs, joinPath(ascendHome, "include")); + std::string driverPath = + getEnvPath("ASCEND_DRIVER_PATH").value_or("/usr/local/Ascend/driver"); + addExistingIncludeDir(includeDirs, joinPath(driverPath, "kernel/inc")); + + if (ptoIsaPath.empty()) { + diagOS << "Warning: PTO_ISA_PATH/PTO_ISA_ROOT is not set; " + "C++ device object emission may fail to include pto/pto-inst.hpp.\n"; + } else if (includeDirs.empty()) { + diagOS << "Warning: no PTO-ISA include directory containing " + "pto/pto-inst.hpp was found under " + << ptoIsaPath << ".\n"; + } + return includeDirs; +} + class VPTOFatobjToolchain; static bool compileDeviceLLVMToObject(llvm::StringRef llPath, @@ -143,7 +204,7 @@ static bool compileDeviceLLVMToObject(llvm::StringRef llPath, llvm::StringRef bishengPath, llvm::StringRef stderrPath, llvm::raw_ostream &diagOS); -static bool compileHostStubToFatobj(llvm::StringRef stubPath, +static bool compileHostStubToObject(llvm::StringRef stubPath, llvm::StringRef outObjPath, llvm::StringRef moduleId, llvm::StringRef targetCPU, @@ -157,8 +218,35 @@ static bool mergeDeviceObjects(llvm::ArrayRef deviceObjPaths, llvm::StringRef stderrPath, llvm::raw_ostream &diagOS); +static llvm::StringRef +getTargetCPU(mlir::pto::ObjectEmissionDeviceTarget target) { + switch (target) { + case mlir::pto::ObjectEmissionDeviceTarget::Vector: + return "dav-c310-vec"; + case mlir::pto::ObjectEmissionDeviceTarget::Cube: + return "dav-c310-cube"; + } + llvm_unreachable("unknown object emission device target"); +} + class VPTOFatobjToolchain { public: + explicit VPTOFatobjToolchain( + const mlir::pto::ObjectEmissionToolchain &toolchain) + : ascendHomePath(toolchain.ascendHomePath), + bishengPath(toolchain.bishengPath), + bishengCc1Path(toolchain.bishengCc1Path), + cceLdPath(toolchain.cceLdPath), + ldLldPath(toolchain.ldLldPath), + resourceDirPath(toolchain.resourceDirPath), + resourceIncludeDirPath(toolchain.resourceIncludeDirPath), + cceStubDirPath(toolchain.cceStubDirPath), + bishengCompilerBinDirPath(toolchain.bishengCompilerBinDirPath), + ptoIsaPath(toolchain.ptoIsaPath) { + cppIncludeDirs.assign(toolchain.cppIncludeDirs.begin(), + toolchain.cppIncludeDirs.end()); + } + static std::optional create(llvm::raw_ostream &diagOS) { std::optional ascendHome = getAscendHomePath(); @@ -167,7 +255,7 @@ class VPTOFatobjToolchain { return std::nullopt; } - VPTOFatobjToolchain toolchain(*ascendHome); + VPTOFatobjToolchain toolchain(*ascendHome, diagOS); if (!toolchain.validate(diagOS)) return std::nullopt; return toolchain; @@ -187,8 +275,26 @@ class VPTOFatobjToolchain { return bishengCompilerBinDirPath; } + mlir::pto::ObjectEmissionToolchain toPublicToolchain() const { + mlir::pto::ObjectEmissionToolchain toolchain; + toolchain.ascendHomePath = ascendHomePath; + toolchain.bishengPath = bishengPath; + toolchain.bishengCc1Path = bishengCc1Path; + toolchain.cceLdPath = cceLdPath; + toolchain.ldLldPath = ldLldPath; + toolchain.resourceDirPath = resourceDirPath; + toolchain.resourceIncludeDirPath = resourceIncludeDirPath; + toolchain.cceStubDirPath = cceStubDirPath; + toolchain.bishengCompilerBinDirPath = bishengCompilerBinDirPath; + toolchain.ptoIsaPath = ptoIsaPath; + toolchain.cppIncludeDirs.assign(cppIncludeDirs.begin(), + cppIncludeDirs.end()); + return toolchain; + } + private: - explicit VPTOFatobjToolchain(llvm::StringRef ascendHome) + explicit VPTOFatobjToolchain(llvm::StringRef ascendHome, + llvm::raw_ostream &diagOS) : ascendHomePath(ascendHome.str()), bishengPath(joinPath(ascendHomePath, "bin/bisheng")), bishengCc1Path( @@ -202,7 +308,10 @@ class VPTOFatobjToolchain { resourceIncludeDirPath(joinPath(resourceDirPath, "include")), cceStubDirPath(joinPath(resourceIncludeDirPath, "cce_stub")), bishengCompilerBinDirPath( - joinPath(ascendHomePath, "tools/bisheng_compiler/bin")) {} + joinPath(ascendHomePath, "tools/bisheng_compiler/bin")) { + cppIncludeDirs = discoverCppIncludeDirs(ascendHomePath, diagOS, + ptoIsaPath); + } bool validate(llvm::raw_ostream &diagOS) const { if (!llvm::sys::fs::exists(bishengPath)) { @@ -234,6 +343,8 @@ class VPTOFatobjToolchain { std::string resourceIncludeDirPath; std::string cceStubDirPath; std::string bishengCompilerBinDirPath; + std::string ptoIsaPath; + llvm::SmallVector cppIncludeDirs; }; class VPTOFatobjArtifacts { @@ -311,7 +422,7 @@ class VPTOFatobjArtifacts { llvm::raw_ostream &diagOS) { if (!tempFiles.create("ptoas-host-stub", ".o", hostStubObjPath, diagOS)) return false; - return compileHostStubToFatobj(stubPath, hostStubObjPath, moduleId, + return compileHostStubToObject(stubPath, hostStubObjPath, moduleId, targetCPU, toolchain, mergedDeviceObjPath, stderrPath, diagOS); } @@ -326,6 +437,9 @@ class VPTOFatobjArtifacts { "-cce-lite-bin-module-id", moduleId.str(), std::string("-cce-aicore-arch=") + targetCPU.str(), + // Keep device externals relocatable until the final + // --cce-fatobj-link link step. + "-dc", "-r", "-o", outPath.str(), @@ -398,6 +512,7 @@ static bool compileDeviceLLVMToObject(llvm::StringRef llPath, std::string("--cce-aicore-arch=") + targetCPU.str(), "--cce-aicore-only", "-O2", + "-dc", "-c", "-x", "ir", @@ -409,7 +524,85 @@ static bool compileDeviceLLVMToObject(llvm::StringRef llPath, "device LLVM compilation", llPath); } -static bool compileHostStubToFatobj(llvm::StringRef stubPath, +static bool compileCppDeviceSourceToObject( + llvm::StringRef cppPath, llvm::StringRef outObjPath, + llvm::StringRef targetCPU, const mlir::pto::ObjectEmissionToolchain &toolchain, + llvm::StringRef stderrPath, llvm::raw_ostream &diagOS) { + llvm::SmallVector args = { + toolchain.bishengPath, + "-xcce", + "-fenable-matrix", + "--cce-aicore-enable-tl", + "--cce-aicore-only", + "-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", + std::string("--cce-aicore-arch=") + targetCPU.str(), + "-DREGISTER_BASE", + "-std=c++17", + "-dc", + }; + for (const std::string &includeDir : toolchain.cppIncludeDirs) + args.push_back("-I" + includeDir); + args.push_back("-c"); + args.push_back(cppPath.str()); + args.push_back("-o"); + args.push_back(outObjPath.str()); + + return runCommandWithStderr(toolchain.bishengPath, args, stderrPath, diagOS, + "C++ device compilation"); +} + +static bool compileCppDeviceSourceToFatobj( + llvm::StringRef cppPath, llvm::StringRef outObjPath, + const mlir::pto::ObjectEmissionToolchain &toolchain, + llvm::StringRef stderrPath, llvm::raw_ostream &diagOS) { + llvm::SmallVector args = { + toolchain.bishengPath, + "-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", + "--cce-aicore-arch=dav-c310", + "-DREGISTER_BASE", + "-std=c++17", + "-O2", + "-dc", + "-c", + }; + for (const std::string &includeDir : toolchain.cppIncludeDirs) + args.push_back("-I" + includeDir); + args.push_back(cppPath.str()); + args.push_back("-o"); + args.push_back(outObjPath.str()); + + return runCommandWithStderr(toolchain.bishengPath, args, stderrPath, diagOS, + "C++ fatobj compilation"); +} + +static bool compileHostStubToObject(llvm::StringRef stubPath, llvm::StringRef outObjPath, llvm::StringRef moduleId, llvm::StringRef targetCPU, @@ -537,13 +730,174 @@ static bool mergeDeviceObjects(llvm::ArrayRef deviceObjPaths, "device object merge"); } +static bool linkFatobjFiles(llvm::ArrayRef fatobjPaths, + llvm::StringRef outObjPath, + const mlir::pto::ObjectEmissionToolchain &toolchain, + llvm::StringRef stderrPath, + llvm::raw_ostream &diagOS) { + if (fatobjPaths.empty()) + return false; + + llvm::SmallVector args = { + toolchain.bishengPath, + "--cce-fatobj-link", + "--cce-aicore-arch=dav-c310", + "-r", + "-o", + outObjPath.str(), + }; + for (const std::string &path : fatobjPaths) + args.push_back(path); + + return runCommandWithStderr(toolchain.bishengPath, args, stderrPath, diagOS, + "fatobj link"); +} + } // namespace -mlir::LogicalResult mlir::pto::emitVPTOFatobj(llvm::Module *cubeModule, - llvm::Module *vectorModule, - llvm::StringRef stubSource, - llvm::ToolOutputFile &outputFile, +mlir::LogicalResult mlir::pto::discoverObjectEmissionToolchain( + ObjectEmissionToolchain &publicToolchain, llvm::raw_ostream &diagOS) { + std::optional toolchain = + VPTOFatobjToolchain::create(diagOS); + if (!toolchain) + return failure(); + publicToolchain = toolchain->toPublicToolchain(); + return success(); +} + +mlir::LogicalResult mlir::pto::writeLLVMModule(llvm::Module &module, + llvm::StringRef path, + llvm::raw_ostream &diagOS) { + return writeLLVMModuleFile(module, path, diagOS) ? success() : failure(); +} + +mlir::LogicalResult mlir::pto::writeCppSource(llvm::StringRef cppSource, + llvm::StringRef path, llvm::raw_ostream &diagOS) { + return writeTextFile(path, cppSource, diagOS) ? success() : failure(); +} + +mlir::LogicalResult mlir::pto::writeHostStubSource( + llvm::StringRef stubSource, llvm::StringRef path, + llvm::raw_ostream &diagOS) { + return writeTextFile(path, stubSource, diagOS) ? success() : failure(); +} + +mlir::LogicalResult mlir::pto::compileCppToDeviceObject( + llvm::StringRef cppPath, llvm::StringRef outObjPath, + ObjectEmissionDeviceTarget target, const ObjectEmissionToolchain &toolchain, + llvm::StringRef stderrPath, llvm::raw_ostream &diagOS) { + return compileCppDeviceSourceToObject(cppPath, outObjPath, + getTargetCPU(target), toolchain, + stderrPath, diagOS) + ? success() + : failure(); +} + +mlir::LogicalResult mlir::pto::compileLLVMToDeviceObject( + llvm::StringRef llPath, llvm::StringRef outObjPath, + ObjectEmissionDeviceTarget target, const ObjectEmissionToolchain &toolchain, + llvm::StringRef stderrPath, llvm::raw_ostream &diagOS) { + return compileDeviceLLVMToObject(llPath, outObjPath, getTargetCPU(target), + toolchain.bishengPath, stderrPath, diagOS) + ? success() + : failure(); +} + +mlir::LogicalResult mlir::pto::emitCppVectorDeviceObject( + llvm::StringRef cppSource, llvm::StringRef cppPath, + llvm::StringRef outObjPath, const ObjectEmissionToolchain &toolchain, + llvm::StringRef stderrPath, llvm::raw_ostream &diagOS) { + if (failed(writeCppSource(cppSource, cppPath, diagOS))) + return failure(); + return compileCppToDeviceObject(cppPath, outObjPath, + ObjectEmissionDeviceTarget::Vector, + toolchain, stderrPath, diagOS); +} + +mlir::LogicalResult mlir::pto::emitCppCubeDeviceObject( + llvm::StringRef cppSource, llvm::StringRef cppPath, + llvm::StringRef outObjPath, const ObjectEmissionToolchain &toolchain, + llvm::StringRef stderrPath, llvm::raw_ostream &diagOS) { + if (failed(writeCppSource(cppSource, cppPath, diagOS))) + return failure(); + return compileCppToDeviceObject(cppPath, outObjPath, + ObjectEmissionDeviceTarget::Cube, + toolchain, stderrPath, diagOS); +} + +mlir::LogicalResult mlir::pto::emitCppFatobj( + llvm::StringRef cppSource, llvm::StringRef cppPath, + llvm::StringRef outObjPath, const ObjectEmissionToolchain &toolchain, + llvm::StringRef stderrPath, llvm::raw_ostream &diagOS) { + if (failed(writeCppSource(cppSource, cppPath, diagOS))) + return failure(); + return compileCppDeviceSourceToFatobj(cppPath, outObjPath, toolchain, + stderrPath, diagOS) + ? success() + : failure(); +} + +mlir::LogicalResult mlir::pto::emitVPTOVectorDeviceObject( + llvm::Module &module, llvm::StringRef llPath, llvm::StringRef outObjPath, + const ObjectEmissionToolchain &toolchain, llvm::StringRef stderrPath, + llvm::raw_ostream &diagOS) { + if (failed(writeLLVMModule(module, llPath, diagOS))) + return failure(); + return compileLLVMToDeviceObject(llPath, outObjPath, + ObjectEmissionDeviceTarget::Vector, + toolchain, stderrPath, diagOS); +} + +mlir::LogicalResult mlir::pto::emitVPTOCubeDeviceObject( + llvm::Module &module, llvm::StringRef llPath, llvm::StringRef outObjPath, + const ObjectEmissionToolchain &toolchain, llvm::StringRef stderrPath, + llvm::raw_ostream &diagOS) { + if (failed(writeLLVMModule(module, llPath, diagOS))) + return failure(); + return compileLLVMToDeviceObject(llPath, outObjPath, + ObjectEmissionDeviceTarget::Cube, + toolchain, stderrPath, diagOS); +} + +mlir::LogicalResult mlir::pto::mergeDeviceObjects( + llvm::ArrayRef deviceObjPaths, llvm::StringRef outObjPath, + const ObjectEmissionToolchain &toolchain, llvm::StringRef stderrPath, + llvm::raw_ostream &diagOS) { + return ::mergeDeviceObjects(deviceObjPaths, outObjPath, toolchain.ldLldPath, + stderrPath, diagOS) + ? success() + : failure(); +} + +mlir::LogicalResult mlir::pto::compileStubToFatobj( + llvm::StringRef stubPath, llvm::StringRef deviceObjPath, + llvm::StringRef outputPath, llvm::StringRef moduleId, + const ObjectEmissionToolchain &toolchain, llvm::StringRef stderrPath, + llvm::raw_ostream &diagOS) { + VPTOFatobjToolchain privateToolchain(toolchain); + + constexpr llvm::StringLiteral targetCPU = "dav-c310"; + return compileHostStubToObject(stubPath, outputPath, moduleId, targetCPU, + privateToolchain, deviceObjPath, stderrPath, + diagOS) + ? success() + : failure(); +} + +mlir::LogicalResult mlir::pto::linkFatobjs( + llvm::ArrayRef fatobjPaths, llvm::StringRef outputPath, + const ObjectEmissionToolchain &toolchain, llvm::StringRef stderrPath, + llvm::raw_ostream &diagOS) { + return linkFatobjFiles(fatobjPaths, outputPath, toolchain, stderrPath, diagOS) + ? success() + : failure(); +} + +mlir::LogicalResult mlir::pto::emitVPTOFatobjFromLLVMModules( + llvm::Module *cubeModule, llvm::Module *vectorModule, + llvm::StringRef stubSource, llvm::ToolOutputFile &outputFile, + llvm::raw_ostream &diagOS) { if (!cubeModule && !vectorModule) { diagOS << "Error: VPTO fatobj emission requires at least one LLVM module.\n"; return failure(); diff --git a/tools/ptoas/ObjectEmission.h b/tools/ptoas/ObjectEmission.h new file mode 100644 index 0000000000..35c8c8f685 --- /dev/null +++ b/tools/ptoas/ObjectEmission.h @@ -0,0 +1,121 @@ +// 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. + +#ifndef PTOAS_OBJECT_EMISSION_H +#define PTOAS_OBJECT_EMISSION_H + +#include "llvm/ADT/StringRef.h" +#include "llvm/ADT/ArrayRef.h" +#include "mlir/Support/LogicalResult.h" + +#include +#include + +namespace llvm { +class ToolOutputFile; +class Module; +class raw_ostream; +} + +namespace mlir::pto { + +enum class ObjectEmissionDeviceTarget { + Vector, + Cube, +}; + +struct ObjectEmissionToolchain { + std::string ascendHomePath; + std::string bishengPath; + std::string bishengCc1Path; + std::string cceLdPath; + std::string ldLldPath; + std::string resourceDirPath; + std::string resourceIncludeDirPath; + std::string cceStubDirPath; + std::string bishengCompilerBinDirPath; + std::string ptoIsaPath; + std::vector cppIncludeDirs; +}; + +LogicalResult discoverObjectEmissionToolchain(ObjectEmissionToolchain &toolchain, + llvm::raw_ostream &diagOS); + +LogicalResult writeLLVMModule(llvm::Module &module, llvm::StringRef path, + llvm::raw_ostream &diagOS); + +LogicalResult writeCppSource(llvm::StringRef cppSource, llvm::StringRef path, + llvm::raw_ostream &diagOS); + +LogicalResult writeHostStubSource(llvm::StringRef stubSource, + llvm::StringRef path, + llvm::raw_ostream &diagOS); + +LogicalResult compileCppToDeviceObject( + llvm::StringRef cppPath, llvm::StringRef outObjPath, + ObjectEmissionDeviceTarget target, const ObjectEmissionToolchain &toolchain, + llvm::StringRef stderrPath, llvm::raw_ostream &diagOS); + +LogicalResult compileLLVMToDeviceObject( + llvm::StringRef llPath, llvm::StringRef outObjPath, + ObjectEmissionDeviceTarget target, const ObjectEmissionToolchain &toolchain, + llvm::StringRef stderrPath, llvm::raw_ostream &diagOS); + +LogicalResult emitCppVectorDeviceObject( + llvm::StringRef cppSource, llvm::StringRef cppPath, + llvm::StringRef outObjPath, const ObjectEmissionToolchain &toolchain, + llvm::StringRef stderrPath, llvm::raw_ostream &diagOS); + +LogicalResult emitCppCubeDeviceObject( + llvm::StringRef cppSource, llvm::StringRef cppPath, + llvm::StringRef outObjPath, const ObjectEmissionToolchain &toolchain, + llvm::StringRef stderrPath, llvm::raw_ostream &diagOS); + +LogicalResult emitCppFatobj(llvm::StringRef cppSource, llvm::StringRef cppPath, + llvm::StringRef outObjPath, + const ObjectEmissionToolchain &toolchain, + llvm::StringRef stderrPath, + llvm::raw_ostream &diagOS); + +LogicalResult emitVPTOVectorDeviceObject( + llvm::Module &module, llvm::StringRef llPath, llvm::StringRef outObjPath, + const ObjectEmissionToolchain &toolchain, llvm::StringRef stderrPath, + llvm::raw_ostream &diagOS); + +LogicalResult emitVPTOCubeDeviceObject( + llvm::Module &module, llvm::StringRef llPath, llvm::StringRef outObjPath, + const ObjectEmissionToolchain &toolchain, llvm::StringRef stderrPath, + llvm::raw_ostream &diagOS); + +LogicalResult mergeDeviceObjects(llvm::ArrayRef deviceObjPaths, + llvm::StringRef outObjPath, + const ObjectEmissionToolchain &toolchain, + llvm::StringRef stderrPath, + llvm::raw_ostream &diagOS); + +LogicalResult compileStubToFatobj( + llvm::StringRef stubPath, llvm::StringRef deviceObjPath, + llvm::StringRef outputPath, llvm::StringRef moduleId, + const ObjectEmissionToolchain &toolchain, llvm::StringRef stderrPath, + llvm::raw_ostream &diagOS); + +LogicalResult linkFatobjs(llvm::ArrayRef fatobjPaths, + llvm::StringRef outputPath, + const ObjectEmissionToolchain &toolchain, + llvm::StringRef stderrPath, + llvm::raw_ostream &diagOS); + +LogicalResult emitVPTOFatobjFromLLVMModules(llvm::Module *cubeModule, + llvm::Module *vectorModule, + llvm::StringRef stubSource, + llvm::ToolOutputFile &outputFile, + llvm::raw_ostream &diagOS); + +} // namespace mlir::pto + +#endif diff --git a/tools/ptoas/PTOASDriver.cpp b/tools/ptoas/PTOASDriver.cpp new file mode 100644 index 0000000000..7ddc7f4a09 --- /dev/null +++ b/tools/ptoas/PTOASDriver.cpp @@ -0,0 +1,2726 @@ +// 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. + +#include "PTOASDriver.h" + +#include "PTO/IR/PTO.h" +#include "PTO/Transforms/VPTOLLVMEmitter.h" +#include "PTO/Transforms/Passes.h" +#include "PTO/Transforms/BufferizableOpInterfaceImpl.h" +#include "ObjectEmission.h" +#include "VPTOHostStubEmission.h" +#include "TilelangDaemon.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/Diagnostics.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/SymbolTable.h" +#include "mlir/InitAllDialects.h" +#include "mlir/InitAllPasses.h" +#include "mlir/Parser/Parser.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Math/IR/Math.h" +#include +#include +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Target/Cpp/CppEmitter.h" +#include "llvm/Support/SourceMgr.h" +#include "llvm/Support/ToolOutputFile.h" +#include "llvm/Support/FileSystem.h" // [Fix] Required for OF_None +#include "llvm/Support/Path.h" +#include "ptobc/ptobc_decode.h" +#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/EmitC/IR/EmitC.h" +#include "mlir/Dialect/EmitC/Transforms/Passes.h" +#include "llvm/Support/CommandLine.h" +#include "llvm/Support/Regex.h" +#include "llvm/Support/raw_ostream.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/ScopeExit.h" +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/ADT/StringSwitch.h" +#include "llvm/ADT/StringMap.h" +#include "llvm/Support/MemoryBuffer.h" +#include "llvm/Support/Process.h" +#include "llvm/Support/Program.h" +#include +#include +#include +#include +#include +#include +#include + +extern "C" { +extern char **environ; +} + +using namespace mlir; +using namespace pto; + +#ifndef PTOAS_RELEASE_VERSION +#define PTOAS_RELEASE_VERSION "unknown" +#endif + +extern "C" int main(int argc, char **argv); + +static void printPTOASVersion(llvm::raw_ostream &os) { + os << "ptoas " << PTOAS_RELEASE_VERSION << "\n"; +} + +static std::string getParentDir(llvm::StringRef path) { + llvm::SmallString<256> parent(path); + llvm::sys::path::remove_filename(parent); + llvm::sys::path::remove_dots(parent, true); + return std::string(parent); +} + +static bool pathExists(llvm::StringRef path) { + return !path.empty() && llvm::sys::fs::exists(path); +} + +static std::string joinPath(llvm::StringRef lhs, llvm::StringRef rhs) { + llvm::SmallString<256> joined(lhs); + llvm::sys::path::append(joined, rhs); + llvm::sys::path::remove_dots(joined, true); + return std::string(joined); +} + +static std::string detectInstalledTilelangPath(const char *argv0) { + std::string exePath = llvm::sys::fs::getMainExecutable(argv0, (void *)&main); + if (exePath.empty()) + return {}; + + const std::string exeDir = getParentDir(exePath); + const std::string prefixDir = getParentDir(exeDir); + const std::string installedTileOps = joinPath(prefixDir, "share/ptoas/TileOps"); + if (pathExists(installedTileOps)) + return installedTileOps; + return {}; +} + +static std::string detectInstalledTilelangPkgPath(const char *argv0) { + std::string exePath = llvm::sys::fs::getMainExecutable(argv0, (void *)&main); + if (exePath.empty()) + return {}; + + const std::string exeDir = getParentDir(exePath); + const std::string prefixDir = getParentDir(exeDir); + const std::string installedPkgRoot = prefixDir; + const std::string installedPkg = joinPath(installedPkgRoot, "tilelang_dsl"); + if (pathExists(installedPkg)) + return installedPkgRoot; + return {}; +} + +static bool hasCLIOption(int argc, char **argv, llvm::StringRef option) { + const std::string optionWithValue = (option + "=").str(); + for (int i = 1; i < argc; ++i) { + llvm::StringRef arg(argv[i]); + if (arg == option || arg.starts_with(optionWithValue)) + return true; + } + return false; +} + +static LogicalResult applyConfiguredPassManagerCLOptions( + PassManager &pm, llvm::StringRef pipelineName, + llvm::raw_ostream &diagOS = llvm::errs()) { + if (succeeded(mlir::applyPassManagerCLOptions(pm))) + return success(); + diagOS << "Error: failed to apply MLIR pass manager command-line options for " + << pipelineName << ".\n"; + return failure(); +} + +static LogicalResult reorderEmitCFunctions(ModuleOp module) { + SmallVector declarations; + SmallVector definitions; + llvm::DenseMap definitionsByName; + + for (auto func : module.getOps()) { + if (func.isDeclaration()) { + declarations.push_back(func); + continue; + } + definitions.push_back(func); + definitionsByName[func.getSymNameAttr()] = func; + } + + llvm::DenseMap indegree; + llvm::DenseMap> outgoing; + for (auto func : definitions) + indegree[func.getOperation()] = 0; + + for (auto caller : definitions) { + Operation *callerOp = caller.getOperation(); + llvm::SmallPtrSet seenCallees; + bool hasCycle = false; + caller.walk([&](emitc::CallOp call) { + auto calleeAttr = call.getCalleeAttr(); + if (!calleeAttr) + return; + auto it = definitionsByName.find(calleeAttr.getLeafReference()); + if (it == definitionsByName.end()) + return; + Operation *calleeOp = it->second.getOperation(); + if (calleeOp == callerOp) { + hasCycle = true; + return; + } + if (!seenCallees.insert(calleeOp).second) + return; + outgoing[calleeOp].push_back(callerOp); + ++indegree[callerOp]; + }); + if (hasCycle) { + return caller.emitOpError() + << "recursive function calls are not supported for EmitC C++ " + "emission"; + } + } + + SmallVector ready; + for (auto func : definitions) { + if (indegree[func.getOperation()] == 0) + ready.push_back(func.getOperation()); + } + + SmallVector sortedDefinitions; + while (!ready.empty()) { + Operation *next = ready.front(); + ready.erase(ready.begin()); + auto nextFunc = cast(next); + sortedDefinitions.push_back(nextFunc); + + for (Operation *user : outgoing[next]) { + unsigned &userIndegree = indegree[user]; + if (--userIndegree == 0) + ready.push_back(user); + } + } + + if (sortedDefinitions.size() != definitions.size()) { + return module.emitError() + << "cyclic function call graph is not supported for EmitC C++ emission"; + } + + if (declarations.empty() && definitions.size() <= 1) + return success(); + + SmallVector desiredOrder; + desiredOrder.append(declarations.begin(), declarations.end()); + desiredOrder.append(sortedDefinitions.begin(), sortedDefinitions.end()); + + Block &body = module.getBodyRegion().front(); + Operation *anchor = nullptr; + for (Operation &op : body.getOperations()) { + if (isa(op)) { + anchor = &op; + break; + } + } + if (!anchor) + return success(); + + auto advanceAnchor = [&]() { + while (anchor) { + anchor = anchor->getNextNode(); + if (!anchor || isa(anchor)) + return; + } + }; + + for (auto func : desiredOrder) { + if (func.getOperation() == anchor) { + advanceAnchor(); + continue; + } + if (anchor) + func->moveBefore(anchor); + else + func->moveBefore(&body, body.end()); + } + + return success(); +} + +// -------------------------------------------------------------------------- +// Command Line Options +// -------------------------------------------------------------------------- +static llvm::cl::opt inputFilename(llvm::cl::Positional, + llvm::cl::desc(""), + llvm::cl::init("-")); + +static llvm::cl::opt outputFilename("o", + llvm::cl::desc("Output filename"), + llvm::cl::value_desc("filename"), + llvm::cl::init("-")); + +static llvm::cl::opt enableInsertSync("enable-insert-sync", + llvm::cl::desc("Enable automatic synchronization insertion pass"), + llvm::cl::init(false)); + +static llvm::cl::opt enableInjectBarrierAllSync( + "enable-inject-barrier-all-sync", + llvm::cl::desc("Enable conservative synchronization by inserting " + "pto.barrier PIPE_ALL before memory-effecting PTO pipe ops"), + llvm::cl::init(false)); + +static llvm::cl::opt enableGraphSyncSolver( + "enable-graph-sync-solver", + llvm::cl::desc("Enable the graph-based intra-core sync solver " + "(experimental). Mutually exclusive with " + "--enable-insert-sync and " + "--enable-inject-barrier-all-sync."), + llvm::cl::init(false)); + +static llvm::cl::opt graphSyncSolverEventIdMax( + "graph-sync-solver-event-id-max", + llvm::cl::desc( + "Maximum EVENT_ID slots for the graph sync solver (default 8). " + "Lower values exercise the PIPE_ALL coloring fallback sooner."), + llvm::cl::init(8)); + +static llvm::cl::opt enableTileOpExpand( + "enable-tile-op-expand", + llvm::cl::desc( + "Deprecated compatibility flag. TileOp expansion is controlled by " + "--pto-backend=vpto."), + llvm::cl::init(false)); + +#ifndef PTOAS_DEFAULT_TILELANG_PATH +#define PTOAS_DEFAULT_TILELANG_PATH "" +#endif +#ifndef PTOAS_DEFAULT_TILELANG_PKG_PATH +#define PTOAS_DEFAULT_TILELANG_PKG_PATH "" +#endif + +static llvm::cl::opt tilelangPath( + "tilelang-path", + llvm::cl::desc("Path to directory of .py tilelang DSL template files " + "(default: /lib/TileOps, baked in at build time)"), + llvm::cl::init(PTOAS_DEFAULT_TILELANG_PATH)); + +static llvm::cl::opt tilelangPkgPath( + "tilelang-pkg-path", + llvm::cl::desc("PYTHONPATH for tilelang_dsl package " + "(default: /tilelang-dsl/python, baked in at build time)"), + llvm::cl::init(PTOAS_DEFAULT_TILELANG_PKG_PATH)); + +static llvm::cl::opt daemonSocketPath( + "daemon-socket-path", + llvm::cl::desc("Path to Unix domain socket for daemon RPC " + "(default: /tmp/tilelang_daemon_{pid}.sock)"), + llvm::cl::init("")); + +static pto::ExpandTileOpOptions resolveExpandTileOpOptions(int argc, + char **argv) { + pto::ExpandTileOpOptions expandOpts; + expandOpts.tilelangPath = tilelangPath; + expandOpts.tilelangPkgPath = tilelangPkgPath; + + if (!hasCLIOption(argc, argv, "--tilelang-path")) { + std::string detectedTilelangPath = detectInstalledTilelangPath(argv[0]); + if (!detectedTilelangPath.empty()) + expandOpts.tilelangPath = detectedTilelangPath; + } + + if (!hasCLIOption(argc, argv, "--tilelang-pkg-path")) { + std::string detectedTilelangPkgPath = detectInstalledTilelangPkgPath(argv[0]); + if (!detectedTilelangPkgPath.empty()) + expandOpts.tilelangPkgPath = detectedTilelangPkgPath; + } + + // Daemon mode is default (no CLI option needed) + // Automatically start daemon for instance caching + if (!expandOpts.tilelangPath.empty()) { + std::string socket = daemonSocketPath; + if (socket.empty()) + socket = ptoas::DaemonManager::generateSocketPath(); + + // Register cleanup handler (daemon will be stopped on PTOAS exit) + ptoas::registerDaemonCleanup(); + + // Try to start daemon automatically + if (ptoas::DaemonManager::start(socket, expandOpts.tilelangPath, expandOpts.tilelangPkgPath)) { + expandOpts.daemonSocketPath = socket; + llvm::errs() << "Info: TileLang daemon started successfully\n"; + } else { + // Fallback: daemon failed, use subprocess mode (current approach) + expandOpts.daemonSocketPath = ""; + llvm::errs() << "Warning: Failed to start daemon, using subprocess mode (fallback)\n"; + } + } + + return expandOpts; +} + +static llvm::cl::opt disableInferLayout( + "disable-infer-layout", + llvm::cl::desc("Disable PTO layout inference pass (static-only)"), + llvm::cl::init(false)); + +static llvm::cl::opt emitAddPtrTrace( + "emit-addptr-trace", + llvm::cl::desc("Emit addptr trace comments in generated C++ output"), + llvm::cl::init(false)); + +static llvm::cl::opt emitMlirIR( + "emit-pto-ir", + llvm::cl::desc("Emit PTO IR after lowering instead of C++"), + llvm::cl::init(false)); + +static llvm::cl::opt ptoTargetArch( + "pto-arch", + llvm::cl::desc("Target Ascend architecture for codegen: a3 or a5 (default: a3)"), + llvm::cl::value_desc("a3|a5"), + llvm::cl::init("a3")); + +static llvm::cl::opt ptoBuildLevel( + "pto-level", + llvm::cl::desc("Build level for pass pipeline: level1, level2, or level3 (default: level2)"), + llvm::cl::value_desc("level1|level2|level3"), + llvm::cl::init("level2")); + +static llvm::cl::opt ptoBackend( + "pto-backend", + llvm::cl::desc("Final PTOAS backend: emitc or vpto (default: emitc)"), + llvm::cl::value_desc("emitc|vpto"), llvm::cl::init("emitc")); + +static llvm::cl::opt emitVPTO( + "emit-vpto", + llvm::cl::desc("Write final post-pass VPTO IR to -o"), + llvm::cl::init(false)); + +static llvm::cl::opt vptoPrintIR( + "vpto-print-ir", + llvm::cl::desc("Print post-pass VPTO backend IR to stderr"), + llvm::cl::init(false)); + +static llvm::cl::opt vptoLoweringStrategy( + "vpto-lowering-strategy", + llvm::cl::desc("VPTO vector lowering strategy: post-update or no-post-update"), + llvm::cl::value_desc("post-update|no-post-update"), + llvm::cl::init("post-update")); + +static llvm::cl::opt dumpVPTOIR( + "dump-vpto-ir", + llvm::cl::desc("Print post-pass VPTO backend IR to stderr"), + llvm::cl::init(false)); + +static llvm::cl::opt ptoPrintSeamIR( + "pto-print-seam-ir", + llvm::cl::desc("Print shared pre-backend seam IR to stderr"), + llvm::cl::init(false)); + +static llvm::cl::opt ptoSeamIRFile( + "pto-seam-ir-file", + llvm::cl::desc("Write shared pre-backend seam IR to a file"), + llvm::cl::value_desc("path"), + llvm::cl::init("")); + +enum class PTOBuildLevel { + Level1, + Level2, + Level3, +}; + +enum class PTOBackend { + EmitC, + VPTO, +}; + +static PTOBuildLevel defaultBuildLevel() { + return PTOBuildLevel::Level2; +} + +static bool parseBuildLevel(llvm::StringRef levelStr, PTOBuildLevel &out) { + std::string s = levelStr.str(); + for (char &c : s) + c = static_cast(std::tolower(static_cast(c))); + if (s == "level1") { + out = PTOBuildLevel::Level1; + return true; + } + if (s == "level2") { + out = PTOBuildLevel::Level2; + return true; + } + if (s == "level3") { + out = PTOBuildLevel::Level3; + return true; + } + return false; +} + +static constexpr llvm::StringLiteral kAutoSyncTailPolicyBarrierAll = + "barrier_all"; +static constexpr llvm::StringLiteral kAutoSyncTailPolicyMte3ToSEvent0 = + "setwait_mte3_to_s_event0"; + +static bool parseAutoSyncTailHint(llvm::StringRef hintStr, std::string &normalized) { + std::string s = hintStr.str(); + for (char &c : s) + c = static_cast(std::tolower(static_cast(c))); + if (s == "barrier-all" || s == "barrier_all" || s == "default") { + normalized = kAutoSyncTailPolicyBarrierAll.str(); + return true; + } + if (s == "mte3-to-s-event0" || s == "mte3_to_s_event0" || + s == "setwait-mte3-to-s-event0" || + s == "setwait_mte3_to_s_event0") { + normalized = kAutoSyncTailPolicyMte3ToSEvent0.str(); + return true; + } + return false; +} + +static bool parseBackend(llvm::StringRef backendStr, PTOBackend &out) { + std::string s = backendStr.str(); + for (char &c : s) + c = static_cast(std::tolower(static_cast(c))); + if (s == "emitc") { + out = PTOBackend::EmitC; + return true; + } + if (s == "vpto") { + out = PTOBackend::VPTO; + return true; + } + return false; +} + +static LogicalResult parsePTOBackendAttr(Operation *op, + std::optional &backend) { + backend = std::nullopt; + Attribute rawBackendAttr = op->getAttr("pto.backend"); + if (!rawBackendAttr) + return success(); + + auto backendAttr = dyn_cast(rawBackendAttr); + if (!backendAttr) { + return op->emitError("invalid pto.backend attribute. Expected string " + "value 'emitc' or 'vpto'."); + } + + PTOBackend attrBackend = PTOBackend::EmitC; + if (!parseBackend(backendAttr.getValue(), attrBackend)) { + return op->emitError("invalid pto.backend '") + << backendAttr.getValue() << "'. Expected 'emitc' or 'vpto'."; + } + + backend = attrBackend; + return success(); +} + +static LogicalResult validatePTOBackendAttrs(ModuleOp module) { + LogicalResult status = success(); + module.walk([&](ModuleOp nested) { + if (failed(status)) + return WalkResult::interrupt(); + std::optional backend; + if (failed(parsePTOBackendAttr(nested.getOperation(), backend))) { + status = failure(); + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + return status; +} + +static LogicalResult verifyBackendChildShape(ModuleOp module, + bool cliBackendSpecified) { + if (cliBackendSpecified || module->hasAttr("pto.backend")) + return success(); + + bool sawEmitC = false; + bool sawVPTO = false; + SmallVector missingBackendChildren; + for (ModuleOp child : module.getOps()) { + std::optional childBackend; + if (failed(parsePTOBackendAttr(child.getOperation(), childBackend))) + return failure(); + if (!childBackend) { + missingBackendChildren.push_back(child); + continue; + } + sawEmitC |= *childBackend == PTOBackend::EmitC; + sawVPTO |= *childBackend == PTOBackend::VPTO; + } + + if (!sawEmitC || !sawVPTO) + return success(); + + if (!missingBackendChildren.empty()) { + return missingBackendChildren.front().emitError() + << "mixed-backend container child module is missing pto.backend"; + } + + return success(); +} + +static LogicalResult resolveModuleBackend(ModuleOp module, + bool cliBackendSpecified, + bool &mixedBackendMode, + PTOBackend &effectiveBackend) { + mixedBackendMode = false; + if (cliBackendSpecified) + return success(); + + std::optional topBackend; + if (failed(parsePTOBackendAttr(module.getOperation(), topBackend))) + return failure(); + if (topBackend) { + effectiveBackend = *topBackend; + return success(); + } + + std::optional childBackend; + bool sawEmitCChild = false; + bool sawVPTOChild = false; + SmallVector missingBackendChildren; + for (ModuleOp child : module.getOps()) { + std::optional parsedChildBackend; + if (failed(parsePTOBackendAttr(child.getOperation(), parsedChildBackend))) + return failure(); + if (!parsedChildBackend) { + missingBackendChildren.push_back(child); + continue; + } + sawEmitCChild |= *parsedChildBackend == PTOBackend::EmitC; + sawVPTOChild |= *parsedChildBackend == PTOBackend::VPTO; + + if (!childBackend) { + childBackend = *parsedChildBackend; + continue; + } + if (*childBackend != *parsedChildBackend) { + if (!missingBackendChildren.empty()) { + return missingBackendChildren.front().emitError() + << "mixed-backend container child module is missing pto.backend"; + } + mixedBackendMode = true; + continue; + } + } + + if (childBackend) + effectiveBackend = *childBackend; + if (sawEmitCChild && sawVPTOChild && !missingBackendChildren.empty()) { + return missingBackendChildren.front().emitError() + << "mixed-backend container child module is missing pto.backend"; + } + return success(); +} + +static bool hasFunctionKernelKind(ModuleOp module) { + return static_cast(module->getAttrOfType( + FunctionKernelKindAttr::name)); +} + +static llvm::StringRef getPublicABISuffix(llvm::StringRef symName) { + if (symName.ends_with(".vector")) + return ".vector"; + if (symName.ends_with(".cube")) + return ".cube"; + return {}; +} + +static llvm::StringRef getDeviceABISuffix(llvm::StringRef symName) { + if (symName.ends_with("_mix_aiv")) + return "_mix_aiv"; + if (symName.ends_with("_mix_aic")) + return "_mix_aic"; + return {}; +} + +static LogicalResult verifyVPTOSourceSymbolNames(ModuleOp module) { + LogicalResult status = success(); + module.walk([&](func::FuncOp func) { + if (failed(status)) + return WalkResult::interrupt(); + + llvm::StringRef symName = func.getSymName(); + if (func->hasAttr("pto.aicore")) { + llvm::StringRef suffix = getDeviceABISuffix(symName); + if (!suffix.empty()) { + status = func.emitError() + << "source-level pto.aicore function '" << symName + << "' must not use reserved device ABI suffix '" << suffix + << "'; PTOAS adds it from pto.kernel_kind"; + return WalkResult::interrupt(); + } + return WalkResult::advance(); + } + + if (!func.isPublic()) + return WalkResult::advance(); + + llvm::StringRef suffix = getPublicABISuffix(symName); + if (!suffix.empty()) { + status = func.emitError() + << "source-level public non-pto.aicore function '" << symName + << "' must not use reserved VPTO ABI suffix '" << suffix + << "'; PTOAS adds it from pto.kernel_kind"; + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + return status; +} + +static LogicalResult verifyVPTOBackendChild(ModuleOp child) { + if (!hasFunctionKernelKind(child)) { + return child.emitError() + << "vpto backend child module must carry 'pto.kernel_kind'"; + } + return verifyVPTOSourceSymbolNames(child); +} + +struct BackendExportPlan { + std::string sourceName; + std::string abiName; + FunctionType type; +}; + +struct BackendImportPlan { + std::string sourceName; + std::string abiName; + FunctionType type; +}; + +struct BackendChildPlan { + ModuleOp sourceModule; + ModuleOp module; + OwningOpRef ownedModule; + PTOBackend backend = PTOBackend::EmitC; + std::optional kernelKind; + SmallVector exports; + SmallVector imports; + SmallVector sourceDeviceSymbols; + SmallVector deviceABISymbols; +}; + +enum class CompilationMode { + EmitC, + VPTO, + MixedFatobj, +}; + +enum class BackendArtifactKind { + VectorObject, + CubeObject, + GenericObject, + Fatobj, +}; + +struct BackendArtifact { + PTOBackend backend = PTOBackend::EmitC; + BackendArtifactKind kind = BackendArtifactKind::GenericObject; + std::string path; + SmallVector exportedABISymbols; +}; + +struct BackendFatobjGroup { + PTOBackend backend = PTOBackend::EmitC; + BackendArtifactKind kind = BackendArtifactKind::GenericObject; + SmallVector objectPaths; +}; + +struct CompilationPlan { + CompilationMode mode = CompilationMode::EmitC; + PTOBackend effectiveBackend = PTOBackend::EmitC; + SmallVector children; +}; + +static std::optional getModuleKernelKind(ModuleOp module) { + auto kernelKind = module->getAttrOfType( + FunctionKernelKindAttr::name); + if (!kernelKind) + return std::nullopt; + return kernelKind.getKernelKind(); +} + +static llvm::StringRef getPublicABISuffix(FunctionKernelKind kind) { + switch (kind) { + case FunctionKernelKind::Vector: + return ".vector"; + case FunctionKernelKind::Cube: + return ".cube"; + } + llvm_unreachable("unknown function kernel kind"); +} + +static llvm::StringRef getDeviceABISuffix(FunctionKernelKind kind) { + switch (kind) { + case FunctionKernelKind::Vector: + return "_mix_aiv"; + case FunctionKernelKind::Cube: + return "_mix_aic"; + } + llvm_unreachable("unknown function kernel kind"); +} + +static void collectBackendChildSymbols(BackendChildPlan &plan) { + for (func::FuncOp func : plan.module.getOps()) { + llvm::StringRef symName = func.getSymName(); + if (func->hasAttr("pto.aicore")) { + plan.sourceDeviceSymbols.push_back(symName.str()); + if (plan.backend == PTOBackend::VPTO && plan.kernelKind) { + plan.deviceABISymbols.push_back( + (symName + getDeviceABISuffix(*plan.kernelKind)).str()); + } + continue; + } + + if (func.isDeclaration()) { + BackendImportPlan import; + import.sourceName = symName.str(); + import.type = func.getFunctionType(); + plan.imports.push_back(std::move(import)); + continue; + } + + if (!func.isPublic()) + continue; + + BackendExportPlan exportPlan; + exportPlan.sourceName = symName.str(); + exportPlan.type = func.getFunctionType(); + if (plan.backend == PTOBackend::VPTO && plan.kernelKind) { + exportPlan.abiName = + (symName + getPublicABISuffix(*plan.kernelKind)).str(); + } else { + exportPlan.abiName = symName.str(); + } + plan.exports.push_back(std::move(exportPlan)); + } +} + +static OwningOpRef cloneBackendChildModule(ModuleOp outer, + ModuleOp child) { + Operation *clonedOp = child.getOperation()->clone(); + auto clonedModule = cast(clonedOp); + + for (NamedAttribute attr : outer->getAttrs()) { + StringRef attrName = attr.getName().getValue(); + if (attrName == SymbolTable::getSymbolAttrName() || attrName == "pto.backend") + continue; + if (!clonedModule->hasAttr(attr.getName())) + clonedModule->setAttr(attr.getName(), attr.getValue()); + } + + return OwningOpRef(clonedModule); +} + +static LogicalResult +collectBackendChildren(ModuleOp module, + SmallVectorImpl &children) { + for (ModuleOp child : module.getOps()) { + std::optional childBackend; + if (failed(parsePTOBackendAttr(child.getOperation(), childBackend))) + return failure(); + if (!childBackend) + continue; + BackendChildPlan plan; + plan.sourceModule = child; + plan.ownedModule = cloneBackendChildModule(module, child); + plan.module = plan.ownedModule.get(); + plan.backend = *childBackend; + plan.kernelKind = getModuleKernelKind(plan.module); + collectBackendChildSymbols(plan); + children.push_back(std::move(plan)); + } + return success(); +} + +static CompilationMode getCompilationMode(PTOBackend effectiveBackend, + ArrayRef children) { + bool sawEmitC = false; + bool sawVPTO = false; + for (const BackendChildPlan &child : children) { + sawEmitC |= child.backend == PTOBackend::EmitC; + sawVPTO |= child.backend == PTOBackend::VPTO; + } + if (sawEmitC && sawVPTO) + return CompilationMode::MixedFatobj; + return effectiveBackend == PTOBackend::VPTO ? CompilationMode::VPTO + : CompilationMode::EmitC; +} + +static LogicalResult verifyUniqueBackendABIExports( + SmallVectorImpl &children) { + llvm::StringMap seen; + for (BackendChildPlan &child : children) { + for (const BackendExportPlan &exportPlan : child.exports) { + auto inserted = seen.try_emplace(exportPlan.abiName, child.module.getLoc()); + if (!inserted.second) { + return child.module.emitError() + << "generated ABI export symbol '" << exportPlan.abiName + << "' is defined by multiple backend child modules"; + } + } + } + return success(); +} + +static LogicalResult verifyExternalImportShape( + SmallVectorImpl &children) { + for (BackendChildPlan &child : children) { + for (func::FuncOp func : child.module.getOps()) { + if (!func.isDeclaration()) + continue; + if (!func.isPrivate()) { + return func.emitError() + << "external func.func declaration must use private visibility"; + } + if (func->hasAttr("pto.aicore")) { + return func.emitError() + << "external func.func declaration must not be marked pto.aicore"; + } + } + } + return success(); +} + +static LogicalResult resolveExternalImports( + SmallVectorImpl &children) { + for (BackendChildPlan &importingChild : children) { + for (BackendImportPlan &importPlan : importingChild.imports) { + const BackendExportPlan *matchedExport = nullptr; + const BackendChildPlan *matchedChild = nullptr; + for (const BackendChildPlan &exportingChild : children) { + if (&exportingChild == &importingChild) + continue; + for (const BackendExportPlan &exportPlan : exportingChild.exports) { + if (exportPlan.sourceName != importPlan.sourceName) + continue; + if (matchedExport) { + return importingChild.module.emitError() + << "cross-backend external import '" << importPlan.sourceName + << "' matches multiple exported definitions"; + } + matchedExport = &exportPlan; + matchedChild = &exportingChild; + } + } + + if (!matchedExport) { + return importingChild.module.emitError() + << "cross-backend external import '" << importPlan.sourceName + << "' has no matching public non-pto.aicore definition in " + "another backend child module"; + } + + if (matchedExport->type != importPlan.type) { + return importingChild.module.emitError() + << "cross-backend external import '" << importPlan.sourceName + << "' signature does not match exported definition"; + } + + importPlan.abiName = matchedExport->abiName; + (void)matchedChild; + } + } + return success(); +} + +static LogicalResult applyResolvedExternalImports(BackendChildPlan &child) { + llvm::StringMap importABIMap; + for (const BackendImportPlan &importPlan : child.imports) { + if (importPlan.abiName.empty()) + return child.module.emitError() + << "unresolved cross-backend external import '" + << importPlan.sourceName << "'"; + importABIMap[importPlan.sourceName] = importPlan.abiName; + } + if (importABIMap.empty()) + return success(); + + for (func::FuncOp func : child.module.getOps()) { + auto it = importABIMap.find(func.getSymName()); + if (it == importABIMap.end()) + continue; + func->setAttr("pto.external_abi", + StringAttr::get(func.getContext(), it->second)); + } + return success(); +} + +static LogicalResult applyVPTOPublicABIExportNames(ModuleOp module) { + LogicalResult status = success(); + module.walk([&](ModuleOp child) { + if (failed(status)) + return WalkResult::interrupt(); + std::optional kernelKind = getModuleKernelKind(child); + if (!kernelKind) + return WalkResult::advance(); + + llvm::StringMap renameMap; + for (func::FuncOp func : child.getOps()) { + if (!func.isPublic() || func.isDeclaration() || + func->hasAttr("pto.aicore")) + continue; + llvm::StringRef oldName = func.getSymName(); + std::string newName = (oldName + getPublicABISuffix(*kernelKind)).str(); + renameMap[oldName] = newName; + } + + for (func::FuncOp func : child.getOps()) { + func.walk([&](func::CallOp call) { + auto it = renameMap.find(call.getCallee()); + if (it != renameMap.end()) + call.setCallee(it->second); + }); + } + + for (func::FuncOp func : child.getOps()) { + auto it = renameMap.find(func.getSymName()); + if (it != renameMap.end()) + func.setSymName(it->second); + } + return WalkResult::advance(); + }); + return status; +} + +static LogicalResult verifyCompilationPlan(ModuleOp module, + bool cliBackendSpecified, + CompilationPlan &plan) { + if (failed(validatePTOBackendAttrs(module))) + return failure(); + if (failed(verifyBackendChildShape(module, cliBackendSpecified))) + return failure(); + + bool checkedVPTOChild = false; + for (ModuleOp child : module.getOps()) { + std::optional childBackend; + if (failed(parsePTOBackendAttr(child.getOperation(), childBackend))) + return failure(); + if (childBackend == PTOBackend::VPTO) { + checkedVPTOChild = true; + if (failed(verifyVPTOBackendChild(child))) + return failure(); + } else if (!childBackend && plan.effectiveBackend == PTOBackend::VPTO && + hasFunctionKernelKind(child)) { + checkedVPTOChild = true; + if (failed(verifyVPTOSourceSymbolNames(child))) + return failure(); + } + } + + if (!checkedVPTOChild && plan.effectiveBackend == PTOBackend::VPTO) + return verifyVPTOSourceSymbolNames(module); + + if (failed(verifyUniqueBackendABIExports(plan.children))) + return failure(); + if (failed(verifyExternalImportShape(plan.children))) + return failure(); + if (failed(resolveExternalImports(plan.children))) + return failure(); + + return success(); +} + +static LogicalResult buildCompilationPlan(ModuleOp module, + bool cliBackendSpecified, + PTOBackend cliOrDefaultBackend, + CompilationPlan &plan) { + plan = CompilationPlan{}; + plan.effectiveBackend = cliOrDefaultBackend; + + bool mixedBackendMode = false; + if (failed(resolveModuleBackend(module, cliBackendSpecified, mixedBackendMode, + plan.effectiveBackend))) + return failure(); + if (failed(collectBackendChildren(module, plan.children))) + return failure(); + plan.mode = mixedBackendMode ? CompilationMode::MixedFatobj + : getCompilationMode(plan.effectiveBackend, + plan.children); + if (failed(verifyCompilationPlan(module, cliBackendSpecified, plan))) + return failure(); + return success(); +} + +static LogicalResult emitSharedPreBackendSeamIR(ModuleOp module, + llvm::StringRef outputPath) { + if (outputPath.empty()) + return success(); + + if (outputPath == "-") { + module->print(llvm::outs()); + llvm::outs() << "\n"; + llvm::outs().flush(); + return success(); + } + + std::error_code ec; + llvm::ToolOutputFile outputFile(outputPath, ec, llvm::sys::fs::OF_None); + if (ec) { + llvm::errs() << "Error: failed to open seam IR file '" << outputPath + << "': " << ec.message() << "\n"; + return failure(); + } + + module->print(outputFile.os()); + outputFile.os() << "\n"; + outputFile.keep(); + return success(); +} + +static bool hasUnexpandedTileOps(ModuleOp module) { + bool found = false; + module.walk([&](Operation *op) { + if (found) + return; + if (isa(op)) + found = true; + }); + return found; +} + +static bool hasTilelangInlineHelpers(ModuleOp module) { + bool found = false; + module.walk([&](func::FuncOp func) { + if (found) + return; + if (func->hasAttr("pto.tilelang.inline_proc")) + found = true; + }); + return found; +} + +// -------------------------------------------------------------------------- +// Post-process C++ output: rewrite marker calls into Tile member calls. +// +// We emit marker calls in EmitC IR because EmitC currently does not provide a +// first-class op for member-function invocation. After translation, we rewrite: +// PTOAS__TILE_SET_VALUE(dst, offset, val) -> dst.SetValue(offset, val) +// PTOAS__TILE_GET_VALUE(src, offset) -> src.GetValue(offset) +// PTOAS__TILE_DATA(obj) -> obj.data() +// PTOAS__TILE_SET_VALIDSHAPE(obj, r, c) -> obj.SetValidShape(r, c) +// PTOAS__PTR_LOAD(ptr, offset) -> ptr[offset] +// PTOAS__PTR_STORE(ptr, offset, val) -> ptr[offset] = val +// PTOAS__EVENTID_ARRAY_LOAD(arr, idx) -> arr[idx] +// PTOAS__EVENTID_ARRAY_STORE(arr, idx, v) -> arr[idx] = v +// -------------------------------------------------------------------------- +struct ParsedMarkerCall { + size_t markerPos = std::string::npos; + size_t rparenPos = std::string::npos; + llvm::SmallVector args; +}; + +struct MarkerRewriteSpec { + llvm::StringRef marker; + llvm::StringRef memberName; + unsigned expectedNumArgs = 0; +}; + +struct MarkerSubscriptRewriteSpec { + llvm::StringRef marker; + unsigned expectedNumArgs = 0; + bool isStore = false; +}; + +static bool parseMarkerArgs(llvm::StringRef argsRef, + llvm::SmallVectorImpl &args) { + size_t partBegin = 0; + int parenDepth = 0; + for (size_t i = 0; i < argsRef.size(); ++i) { + char c = argsRef[i]; + if (c == '(') { + ++parenDepth; + continue; + } + if (c == ')') { + if (parenDepth > 0) + --parenDepth; + continue; + } + if (c == ',' && parenDepth == 0) { + args.push_back(argsRef.slice(partBegin, i).trim()); + partBegin = i + 1; + } + } + if (partBegin > argsRef.size()) + return false; + args.push_back(argsRef.drop_front(partBegin).trim()); + return true; +} + +static std::optional +findNextMarkerCall(const std::string &cpp, llvm::StringRef marker, + size_t searchPos) { + ParsedMarkerCall call; + call.markerPos = cpp.find(marker.str(), searchPos); + if (call.markerPos == std::string::npos) + return std::nullopt; + + size_t lparenPos = call.markerPos + marker.size(); + if (lparenPos >= cpp.size() || cpp[lparenPos] != '(') + return ParsedMarkerCall{call.markerPos, std::string::npos, {}}; + + size_t argsBegin = lparenPos + 1; + int parenDepth = 0; + for (size_t i = argsBegin; i < cpp.size(); ++i) { + char c = cpp[i]; + if (c == '(') { + ++parenDepth; + continue; + } + if (c != ')') + continue; + if (parenDepth == 0) { + call.rparenPos = i; + break; + } + --parenDepth; + } + if (call.rparenPos == std::string::npos) + return call; + + llvm::StringRef argsRef(cpp.data() + argsBegin, call.rparenPos - argsBegin); + if (!parseMarkerArgs(argsRef, call.args)) + call.args.clear(); + return call; +} + +template +static bool rewriteMarkerCalls(std::string &cpp, llvm::StringRef marker, + BuildReplacementFn buildReplacement) { + size_t searchPos = 0; + bool changed = false; + for (auto call = findNextMarkerCall(cpp, marker, searchPos); call; + call = findNextMarkerCall(cpp, marker, searchPos)) { + if (call->rparenPos == std::string::npos) { + searchPos = call->markerPos + marker.size(); + continue; + } + + std::optional replacement = buildReplacement(*call); + if (!replacement) { + searchPos = call->rparenPos + 1; + continue; + } + + cpp.replace(call->markerPos, (call->rparenPos - call->markerPos) + 1, + *replacement); + changed = true; + searchPos = call->markerPos + replacement->size(); + } + return changed; +} + +static bool rewriteMarkerCallToMember(std::string &cpp, llvm::StringRef marker, + llvm::StringRef memberName, + unsigned expectedNumArgs) { + return rewriteMarkerCalls( + cpp, marker, [&](const ParsedMarkerCall &call) -> std::optional { + if (call.args.size() != expectedNumArgs) + return std::nullopt; + + std::string replacement; + replacement.reserve(marker.size() + 16); + replacement.append(call.args[0].str()); + replacement.push_back('.'); + replacement.append(memberName.str()); + replacement.push_back('('); + if (expectedNumArgs >= 2) + replacement.append(call.args[1].str()); + if (expectedNumArgs == 3) { + replacement.append(", "); + replacement.append(call.args[2].str()); + } + replacement.push_back(')'); + return replacement; + }); +} + +static void rewriteMarkerCallsToMembers( + std::string &cpp, llvm::ArrayRef rewrites) { + bool changed = true; + while (changed) { + changed = false; + for (const MarkerRewriteSpec &rewrite : rewrites) { + changed |= rewriteMarkerCallToMember(cpp, rewrite.marker, + rewrite.memberName, + rewrite.expectedNumArgs); + } + } +} + +static void rewriteTileGetSetValueMarkers(std::string &cpp) { + static const MarkerRewriteSpec kTileMarkerRewrites[] = { + {"PTOAS__TILE_SET_VALUE", "SetValue", 3}, + {"PTOAS__TILE_GET_VALUE", "GetValue", 2}, + {"PTOAS__TILE_DATA", "data", 1}, + {"PTOAS__TILE_SET_VALIDSHAPE", "SetValidShape", 3}, + }; + rewriteMarkerCallsToMembers(cpp, kTileMarkerRewrites); +} + +static void rewriteAsyncEventMarkers(std::string &cpp) { + static const MarkerRewriteSpec kAsyncEventMarkerRewrites[] = { + {"PTOAS__ASYNC_EVENT_WAIT", "Wait", 2}, + {"PTOAS__ASYNC_EVENT_TEST", "Test", 2}, + }; + rewriteMarkerCallsToMembers(cpp, kAsyncEventMarkerRewrites); +} + +// -------------------------------------------------------------------------- +// EmitC cleanup: drop empty emitc.expression ops. +// +// After FormExpressions + CSE, EmitC expressions can become empty when their +// root op is CSE'd with an equivalent dominating value outside the expression +// region. Such expressions crash mlir::emitc::translateToCpp because +// ExpressionOp::getRootOp() returns nullptr. +// -------------------------------------------------------------------------- +static void dropEmptyEmitCExpressions(Operation *rootOp) { + llvm::SmallVector toErase; + rootOp->walk([&](emitc::ExpressionOp expr) { + if (expr.getRootOp()) + return; + Block *body = expr.getBody(); + if (!body) + return; + auto yield = dyn_cast(body->getTerminator()); + if (!yield || yield.getNumOperands() != 1) + return; + Value yielded = yield.getOperand(0); + expr.getResult().replaceAllUsesWith(yielded); + toErase.push_back(expr); + }); + for (emitc::ExpressionOp expr : llvm::reverse(toErase)) + expr.erase(); +} + +static Attribute getDefaultEmitCVariableInitAttr(OpBuilder &builder, Type type) { + if (auto intTy = dyn_cast(type)) + return builder.getIntegerAttr(intTy, 0); + if (isa(type)) + return builder.getIndexAttr(0); + if (auto floatTy = dyn_cast(type)) + return builder.getFloatAttr(floatTy, 0.0); + if (isa(type)) + return emitc::OpaqueAttr::get(builder.getContext(), ""); + return Attribute{}; +} + +// FormExpressions may inline conditions into emitc.expression, but the C++ +// emitter prints cf.br/cf.cond_br operands by variable name rather than by +// recursively emitting an expression. Materialize such operands so CFG-based +// lowering (e.g. scf.while -> cf.*) stays valid. +static void materializeControlFlowOperands(Operation *rootOp) { + llvm::SmallVector branches; + rootOp->walk([&](Operation *op) { + if (isa(op)) + branches.push_back(op); + }); + + OpBuilder builder(rootOp->getContext()); + for (Operation *op : branches) { + builder.setInsertionPoint(op); + for (OpOperand &operand : op->getOpOperands()) { + Value value = operand.get(); + auto expr = dyn_cast_or_null(value.getDefiningOp()); + if (!expr) + continue; + + Attribute initAttr = + getDefaultEmitCVariableInitAttr(builder, value.getType()); + if (!initAttr) + continue; + + Value tmp = + builder.create(op->getLoc(), value.getType(), + initAttr) + .getResult(); + builder.create(op->getLoc(), tmp, value); + operand.set(tmp); + } + } +} + +static bool rewriteMarkerCallToSubscript(std::string &cpp, llvm::StringRef marker, + unsigned expectedNumArgs, + bool isStore) { + return rewriteMarkerCalls( + cpp, marker, [&](const ParsedMarkerCall &call) -> std::optional { + if (call.args.size() != expectedNumArgs) + return std::nullopt; + if (isStore) { + return (call.args[0] + "[" + call.args[1] + "] = " + call.args[2]) + .str(); + } + return (call.args[0] + "[" + call.args[1] + "]").str(); + }); +} + +static void rewriteMarkerCallsToSubscripts( + std::string &cpp, llvm::ArrayRef rewrites) { + bool changed = true; + while (changed) { + changed = false; + for (const MarkerSubscriptRewriteSpec &rewrite : rewrites) { + changed |= rewriteMarkerCallToSubscript(cpp, rewrite.marker, + rewrite.expectedNumArgs, + rewrite.isStore); + } + } +} + +static void rewritePtrScalarMarkers(std::string &cpp) { + static const MarkerSubscriptRewriteSpec kPtrMarkerRewrites[] = { + {"PTOAS__PTR_LOAD", 2, false}, + {"PTOAS__PTR_STORE", 3, true}, + }; + rewriteMarkerCallsToSubscripts(cpp, kPtrMarkerRewrites); +} + +static void rewriteEventIdArrayMarkers(std::string &cpp) { + static const MarkerSubscriptRewriteSpec kEventIdMarkerRewrites[] = { + {"PTOAS__EVENTID_ARRAY_LOAD", 2, false}, + {"PTOAS__EVENTID_ARRAY_STORE", 3, true}, + }; + rewriteMarkerCallsToSubscripts(cpp, kEventIdMarkerRewrites); +} + +static bool rewriteAddPtrTraceMarkers(std::string &cpp, bool showTrace) { + size_t searchPos = 0; + bool changed = false; + for (auto call = findNextMarkerCall(cpp, "PTOAS__ADDPTR_TRACE", searchPos); + call; call = findNextMarkerCall(cpp, "PTOAS__ADDPTR_TRACE", searchPos)) { + if (call->rparenPos == std::string::npos) { + searchPos = call->markerPos + 1; + continue; + } + if (call->args.size() != 3) { + searchPos = call->rparenPos + 1; + continue; + } + + std::string replacement; + if (showTrace) { + replacement.reserve(64); + replacement.append("/* ADDPTR_TRACE: "); + replacement.append(call->args[0].str()); + replacement.append(" = "); + replacement.append(call->args[1].str()); + replacement.append(" + "); + replacement.append(call->args[2].str()); + replacement.append(" */"); + } + + size_t replaceEnd = call->rparenPos; + if (!showTrace) { + size_t i = call->rparenPos + 1; + while (i < cpp.size() && std::isspace(static_cast(cpp[i]))) + ++i; + if (i < cpp.size() && cpp[i] == ';') + replaceEnd = i; + } + + cpp.replace(call->markerPos, (replaceEnd - call->markerPos) + 1, + replacement); + changed = true; + searchPos = call->markerPos + replacement.size(); + } + return changed; +} + +static bool isGeneratedGlobalTensorDecl(llvm::StringRef trimmed, + llvm::StringRef &decl, + llvm::StringRef &varName) { + if (!trimmed.starts_with("GlobalTensor<") || !trimmed.ends_with(";") || + trimmed.contains('=') || trimmed.contains('(')) { + return false; + } + + decl = trimmed.drop_back().rtrim(); + size_t lastWs = decl.find_last_of(" \t"); + if (lastWs == llvm::StringRef::npos) + return false; + varName = decl.drop_front(lastWs + 1); + if (!varName.starts_with("v") || varName.size() <= 1) + return false; + return llvm::all_of(varName.drop_front(1), + [](char c) { return std::isdigit(c); }); +} + +static void rewriteHoistedGlobalTensorDecls(std::string &cpp) { + // When `declareVariablesAtTop` is enabled, the C++ emitter hoists SSA value + // declarations to the top of the function and emits assignments later. This + // requires the C++ type to be default-constructible. + // + // `GlobalTensor<...>` from pto-isa does NOT have a default constructor, so a + // hoisted declaration like: + // GlobalTensor<...> v42; + // fails to compile. Initialize those hoisted temporaries with a null pointer + // so they are constructible: + // GlobalTensor<...> v42(nullptr); + // + // We keep the assignment later; the null-initialized value is never used. + std::string out; + out.reserve(cpp.size() + 64); + + llvm::StringRef ref(cpp); + while (!ref.empty()) { + auto split = ref.split('\n'); + llvm::StringRef line = split.first; + llvm::StringRef rest = split.second; + + llvm::StringRef trimmed = line.trim(); + bool rewritten = false; + llvm::StringRef decl; + llvm::StringRef varName; + if (isGeneratedGlobalTensorDecl(trimmed, decl, varName)) { + size_t indentLen = line.find_first_not_of(" \t"); + if (indentLen == std::string::npos) + indentLen = 0; + llvm::StringRef indent = line.take_front(indentLen); + + out.append(indent.str()); + out.append(decl.str()); + out.append("(nullptr);"); + rewritten = true; + } + + if (!rewritten) + out.append(line.str()); + if (!rest.empty()) + out.push_back('\n'); + ref = rest; + } + + cpp.swap(out); +} + +namespace { +struct ConstantDeclCandidate { + size_t declLine = 0; + std::string indent; + std::string type; + bool hasInitializer = false; + std::string initializer; + size_t assignmentCount = 0; + size_t assignmentLine = 0; + std::string assignmentRhs; +}; +} // namespace + +static bool isGeneratedValueName(llvm::StringRef name) { + if (!name.consume_front("v") || name.empty()) + return false; + return llvm::all_of(name, [](char c) { return std::isdigit(c); }); +} + +static bool isConstFoldableScalarType(llvm::StringRef type) { + type = type.trim(); + if (type.starts_with("const ") || type.starts_with("constexpr ")) + return false; + return llvm::StringSwitch(type) + .Cases("bool", "float", "double", "half", "bfloat16_t", true) + .Cases("int8_t", "uint8_t", "int16_t", "uint16_t", true) + .Cases("int32_t", "uint32_t", "int64_t", "uint64_t", true) + .Default(false); +} + +static bool isLiteralInitializer(llvm::StringRef rhs) { + rhs = rhs.trim(); + if (rhs.empty()) + return false; + if (rhs == "true" || rhs == "false" || rhs == "nullptr") + return true; + + static const llvm::Regex kIntLiteral( + R"(^[+-]?(0[xX][0-9A-Fa-f]+|[0-9]+)[uUlL]*$)"); + static const llvm::Regex kFloatLiteral( + R"(^[+-]?(([0-9]+\.[0-9]*|\.[0-9]+|[0-9]+)([eE][+-]?[0-9]+)?|[0-9]+[eE][+-]?[0-9]+)[fF]?$)"); + static const llvm::Regex kHexFloatLiteral( + R"(^[+-]?0[xX]([0-9A-Fa-f]+\.[0-9A-Fa-f]*|[0-9A-Fa-f]+|\.[0-9A-Fa-f]+)[pP][+-]?[0-9]+[fF]?$)"); + static const llvm::Regex kSpecialFloatLiteral( + R"(^[+-]?(nan|inf)[fF]?$)"); + + return kIntLiteral.match(rhs) || kFloatLiteral.match(rhs) || + kHexFloatLiteral.match(rhs) || kSpecialFloatLiteral.match(rhs); +} + +static std::string normalizeConstInitializer(llvm::StringRef type, + llvm::StringRef rhs) { + type = type.trim(); + rhs = rhs.trim(); + if (type == "bool") { + if (rhs == "0" || rhs == "false") + return "false"; + if (rhs == "1" || rhs == "-1" || rhs == "true") + return "true"; + } + return rhs.str(); +} + +static bool parseConstantDeclarationLine(llvm::StringRef line, + ConstantDeclCandidate &candidate, + std::string &valueName) { + llvm::StringRef trimmed = line.trim(); + if (trimmed.empty() || trimmed.starts_with("#") || trimmed.starts_with("//") || + !trimmed.ends_with(";")) + return false; + + llvm::StringRef body = trimmed.drop_back().rtrim(); + if (body.starts_with("return") || body.starts_with("goto ") || + body.starts_with("if ") || body.starts_with("if(") || + body.starts_with("switch ") || body.starts_with("switch(") || + body.starts_with("for ") || body.starts_with("for(") || + body.starts_with("while ") || body.starts_with("while(") || + body.starts_with("case ") || body == "default") + return false; + + llvm::StringRef lhs = body; + llvm::StringRef rhs; + if (size_t eqPos = body.find('='); eqPos != llvm::StringRef::npos) { + lhs = body.take_front(eqPos).rtrim(); + rhs = body.drop_front(eqPos + 1).trim(); + } + + size_t lastWs = lhs.find_last_of(" \t"); + if (lastWs == llvm::StringRef::npos) + return false; + + llvm::StringRef type = lhs.take_front(lastWs).rtrim(); + llvm::StringRef name = lhs.drop_front(lastWs + 1).trim(); + if (!isGeneratedValueName(name) || !isConstFoldableScalarType(type)) + return false; + + size_t indentLen = line.find_first_not_of(" \t"); + if (indentLen == llvm::StringRef::npos) + indentLen = 0; + candidate.indent = line.take_front(indentLen).str(); + candidate.type = type.str(); + valueName = name.str(); + + if (!rhs.empty()) { + if (!isLiteralInitializer(rhs)) + return false; + candidate.hasInitializer = true; + candidate.initializer = normalizeConstInitializer(type, rhs); + } + + return true; +} + +static bool parseGeneratedValueAssignment(llvm::StringRef line, + llvm::StringRef &valueName, + llvm::StringRef &rhs) { + llvm::StringRef trimmed = line.trim(); + if (trimmed.empty() || trimmed.starts_with("#") || trimmed.starts_with("//") || + !trimmed.ends_with(";")) + return false; + + llvm::StringRef body = trimmed.drop_back().rtrim(); + size_t eqPos = body.find('='); + if (eqPos == llvm::StringRef::npos) + return false; + + llvm::StringRef lhs = body.take_front(eqPos).rtrim(); + rhs = body.drop_front(eqPos + 1).trim(); + if (!isGeneratedValueName(lhs)) + return false; + valueName = lhs; + return true; +} + +static void rewriteScalarConstantDecls(std::string &cpp) { + llvm::SmallVector lines; + for (llvm::StringRef ref(cpp); !ref.empty(); ref = ref.split('\n').second) { + auto split = ref.split('\n'); + lines.push_back(split.first.str()); + } + + llvm::SmallVector eraseLine(lines.size(), false); + auto rewriteSegment = [&](size_t beginLine, size_t endLine) { + llvm::StringMap candidates; + + for (size_t i = beginLine; i <= endLine; ++i) { + ConstantDeclCandidate candidate; + std::string valueName; + if (parseConstantDeclarationLine(lines[i], candidate, valueName)) { + candidate.declLine = i; + candidates[valueName] = std::move(candidate); + continue; + } + + llvm::StringRef assignedName; + llvm::StringRef rhs; + if (!parseGeneratedValueAssignment(lines[i], assignedName, rhs)) + continue; + + auto it = candidates.find(assignedName); + if (it == candidates.end()) + continue; + + ConstantDeclCandidate &info = it->second; + ++info.assignmentCount; + info.assignmentLine = i; + info.assignmentRhs = rhs.str(); + } + + for (auto &entry : candidates) { + llvm::StringRef valueName = entry.getKey(); + ConstantDeclCandidate &info = entry.getValue(); + + std::string initializer; + if (info.hasInitializer) { + if (info.assignmentCount != 0) + continue; + initializer = info.initializer; + } else { + if (info.assignmentCount != 1) + continue; + if (!isLiteralInitializer(info.assignmentRhs)) + continue; + initializer = normalizeConstInitializer( + info.type, llvm::StringRef(info.assignmentRhs)); + eraseLine[info.assignmentLine] = true; + } + + lines[info.declLine] = (info.indent + "const " + info.type + " " + + valueName.str() + " = " + initializer + ";"); + } + }; + + int braceDepth = 0; + size_t segmentStart = 0; + for (size_t i = 0; i < lines.size(); ++i) { + int depthBefore = braceDepth; + for (char c : lines[i]) { + if (c == '{') + ++braceDepth; + else if (c == '}') + --braceDepth; + } + + if (depthBefore == 0 && braceDepth > 0) + segmentStart = i; + if (depthBefore > 0 && braceDepth == 0) + rewriteSegment(segmentStart, i); + } + + std::string out; + out.reserve(cpp.size()); + for (size_t i = 0; i < lines.size(); ++i) { + if (eraseLine[i]) + continue; + out.append(lines[i]); + if (i + 1 != lines.size()) + out.push_back('\n'); + } + cpp.swap(out); +} + +static bool shouldDeclareVariablesAtTop(ModuleOp module) { + auto hasMultiBlockFunc = [](auto func) { return func.getBlocks().size() > 1; }; + return llvm::any_of(module.getOps(), hasMultiBlockFunc) || + llvm::any_of(module.getOps(), hasMultiBlockFunc); +} + +static LogicalResult emitEmitCSourceArtifact(ModuleOp module, + bool emitAddPtrTraceComments, + std::string &cppOutput, + llvm::raw_ostream &diagOS) { + dropEmptyEmitCExpressions(module); + materializeControlFlowOperands(module); + if (failed(reorderEmitCFunctions(module))) { + diagOS << "Error: Failed to order emitted functions for C++ emission.\n"; + return failure(); + } + + cppOutput.clear(); + llvm::raw_string_ostream cppOS(cppOutput); + bool declareVariablesAtTop = shouldDeclareVariablesAtTop(module); + if (failed(emitc::translateToCpp( + module, cppOS, /*declareVariablesAtTop=*/declareVariablesAtTop))) { + diagOS << "Error: Failed to emit C++.\n"; + return failure(); + } + cppOS.flush(); + rewriteTileGetSetValueMarkers(cppOutput); + rewriteAsyncEventMarkers(cppOutput); + rewritePtrScalarMarkers(cppOutput); + rewriteEventIdArrayMarkers(cppOutput); + rewriteAddPtrTraceMarkers(cppOutput, emitAddPtrTraceComments); + rewriteScalarConstantDecls(cppOutput); + rewriteHoistedGlobalTensorDecls(cppOutput); + return success(); +} + +static LogicalResult configureSharedPTOPipeline(PassManager &pm, + PTOBuildLevel effectiveLevel) { + if (failed(applyPassManagerCLOptions(pm))) + return failure(); + + pm.addNestedPass( + pto::createPTOAssignDefaultFrontendPipeIdPass()); + pm.addNestedPass( + pto::createPTOLowerFrontendPipeOpsPass()); + // pm.addNestedPass(pto::createPTOVerifyTFreePass()); + pm.addPass(pto::createPTOInferValidatePipeInitPass()); + pm.addNestedPass(pto::createLoweringSyncToPipePass()); + + if (!disableInferLayout) + pm.addNestedPass(pto::createInferPTOLayoutPass()); + pm.addNestedPass(pto::createPTOA5NormalizeTMovPass()); + pm.addPass(pto::createPTOViewToMemrefPass()); + + if (effectiveLevel != PTOBuildLevel::Level3) { + PlanMemoryOptions planMemoryOption; + planMemoryOption.memMode = MemPlanMode::LOCAL_MEM_PLAN; + planMemoryOption.enableGlobalReuse = false; + planMemoryOption.enablePrintMemoryAllocatedSize = false; + pm.addPass(pto::createPlanMemoryPass(planMemoryOption)); + } + pm.addPass(pto::createPTOResolveReservedBuffersPass()); + + // Conditionally add one automatic synchronization mode. Barrier-all is a + // conservative standalone pass; InsertSync and GraphSyncSolver are set/wait + // solvers. + if (enableInsertSync) + pm.addNestedPass(pto::createPTOInsertSyncPass()); + else if (enableInjectBarrierAllSync) + pm.addNestedPass( + pto::createPTOInjectBarrierAllSyncPass()); + else if (enableGraphSyncSolver) { + PTOGraphSyncSolverOptions graphSyncOpts; + graphSyncOpts.eventIdNumMax = graphSyncSolverEventIdMax; + pm.addNestedPass( + pto::createPTOGraphSyncSolverPass(graphSyncOpts)); + } + + return success(); +} + +static void addSharedPTOSeamPasses(PassManager &pm) { + // Reintroduce tile-native handles once on the shared mainline so both + // backends consume the same post-planning seam IR. + pm.addPass(pto::createPTOMaterializeTileHandlesPass()); + pm.addPass(createCSEPass()); +} + +static void prepareVPTOForEmission(PassManager &pm) { + auto &kernelModulePM = pm.nest(); + kernelModulePM.addPass(createCanonicalizerPass()); + kernelModulePM.addPass(createCSEPass()); + kernelModulePM.addPass(pto::createVPTOPtrNormalizePass()); + kernelModulePM.addPass(pto::createVPTOPtrCastCleanupPass()); + kernelModulePM.addPass(createReconcileUnrealizedCastsPass()); + kernelModulePM.addNestedPass( + createVPTOExpandWrapperOpsPass()); + kernelModulePM.addPass(createCSEPass()); + kernelModulePM.addNestedPass( + pto::createPTOInferVPTOVecScopePass()); + kernelModulePM.addPass(createCanonicalizerPass()); + kernelModulePM.addPass(createCSEPass()); + kernelModulePM.addPass(pto::createPTOValidateVPTOEmissionIRPass()); +} + +static void lowerPTOToVPTOBackend(PassManager &pm, int argc, char **argv) { + // TileOp Expand path: + // 1. ExpandTileOp: instantiate TileLang DSL templates, replace tile ops + // with func.call to template functions (tile_buf params) + // 2. InlineLibCall: inline template function bodies + // 3. FoldTileBufIntrinsics: fold tile_buf_addr / tile_valid_rows / + // tile_valid_cols to concrete memref/constant values + auto &kernelModulePM = pm.nest(); + pto::ExpandTileOpOptions expandOpts = resolveExpandTileOpOptions(argc, argv); + kernelModulePM.addPass(pto::createExpandTileOpPass(expandOpts)); + + kernelModulePM.addPass(pto::createPTOInlineLibCallPass()); + kernelModulePM.addNestedPass( + pto::createFoldTileBufIntrinsicsPass()); + // FoldTileBufIntrinsics materializes many constant branch conditions. + // Clean them up immediately on the TileOp expansion path before the + // authoring-stage VPTO verifier and let the existing CSE passes remove the + // resulting dead values later in the pipeline. + kernelModulePM.addPass(mlir::createSCCPPass()); + kernelModulePM.addPass(mlir::createCanonicalizerPass()); +} + +static void inlineTilelangHelpersOnVPTOInput(PassManager &pm) { + auto &kernelModulePM = pm.nest(); + kernelModulePM.addPass(pto::createPTOInlineLibCallPass()); + kernelModulePM.addPass(mlir::createSCCPPass()); + kernelModulePM.addPass(mlir::createCanonicalizerPass()); +} + +static pto::VPTOEmissionOptions buildVPTOEmissionOptions() { + pto::VPTOEmissionOptions options; + options.dumpVPTOIR = false; + options.targetTriple = "hiipu64-hisilicon-cce"; + return options; +} + +struct CompiledBackendChild { + const BackendChildPlan *plan = nullptr; + PTOBackend backend = PTOBackend::EmitC; + std::optional kernelKind; + ModuleOp vptoModule; + std::string cppSource; + pto::EmittedLLVMModule cubeModule; + pto::EmittedLLVMModule vectorModule; +}; + +static LogicalResult runVPTOBackendPipeline(ModuleOp module, int argc, + char **argv, + bool hasTileOpsToExpand, + bool hasTilelangHelpers); + +static LogicalResult runSharedPTOToSeamPipeline(ModuleOp module, + PTOBuildLevel effectiveLevel, + llvm::StringRef arch, + llvm::raw_ostream &diagOS) { + PassManager pm(module.getContext()); + if (failed(configureSharedPTOPipeline(pm, effectiveLevel))) + return failure(); + addSharedPTOSeamPasses(pm); + if (failed(applyConfiguredPassManagerCLOptions(pm, "mixed child PTOAS pipeline"))) + return failure(); + module->setAttr("pto.target_arch", StringAttr::get(module.getContext(), arch)); + if (failed(pm.run(module))) { + diagOS << "Error: mixed child PTO pass execution failed.\n"; + return failure(); + } + return success(); +} + +static LogicalResult reconcileExternalDeclarationTypes(ModuleOp module) { + SymbolTable symbolTable(module); + llvm::StringMap> operandTypesByExternalCallee; + LogicalResult status = success(); + + module.walk([&](func::CallOp call) { + if (failed(status)) + return WalkResult::interrupt(); + auto callee = symbolTable.lookup(call.getCallee()); + if (!callee || !callee.isDeclaration()) + return WalkResult::advance(); + + SmallVector operandTypes(call.getOperandTypes().begin(), + call.getOperandTypes().end()); + auto [it, inserted] = + operandTypesByExternalCallee.try_emplace(call.getCallee(), operandTypes); + if (!inserted && it->second != operandTypes) { + status = call.emitError() + << "external declaration '" << call.getCallee() + << "' is called with inconsistent operand types after PTO " + "normalization"; + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + if (failed(status)) + return failure(); + + for (func::FuncOp func : module.getOps()) { + if (!func.isDeclaration()) + continue; + auto it = operandTypesByExternalCallee.find(func.getSymName()); + if (it == operandTypesByExternalCallee.end()) + continue; + if (func.getResultTypes().empty() && + llvm::equal(func.getArgumentTypes(), it->second)) + continue; + if (!func.getResultTypes().empty()) { + return func.emitError() + << "external declarations with results are not supported in mixed " + "backend object emission"; + } + func.setFunctionType( + FunctionType::get(func.getContext(), it->second, TypeRange{})); + } + return success(); +} + +static LogicalResult runEmitCBackendPipeline(ModuleOp module, + PTOBuildLevel effectiveLevel, + llvm::StringRef arch, + llvm::raw_ostream &diagOS) { + PassManager sharedPM(module.getContext()); + // Shared PTO normalization can rewrite pto.aicore argument types before the + // private external declarations that call them are reconciled below. + sharedPM.enableVerifier(false); + if (failed(configureSharedPTOPipeline(sharedPM, effectiveLevel))) + return failure(); + addSharedPTOSeamPasses(sharedPM); + if (failed(applyConfiguredPassManagerCLOptions(sharedPM, + "mixed child shared pipeline"))) + return failure(); + module->setAttr("pto.target_arch", StringAttr::get(module.getContext(), arch)); + if (failed(sharedPM.run(module))) { + diagOS << "Error: mixed EmitC child shared pass execution failed.\n"; + return failure(); + } + + if (failed(reconcileExternalDeclarationTypes(module))) + return failure(); + + PassManager emitCPM(module.getContext()); + if (arch == "a3") + emitCPM.addPass(pto::createEmitPTOManualPass(pto::PTOArch::A3)); + else + emitCPM.addPass(pto::createEmitPTOManualPass(pto::PTOArch::A5)); + emitCPM.addPass(emitc::createFormExpressionsPass()); + emitCPM.addPass(mlir::createCSEPass()); + if (failed(applyConfiguredPassManagerCLOptions(emitCPM, + "mixed child EmitC pipeline"))) + return failure(); + if (failed(emitCPM.run(module))) { + diagOS << "Error: mixed EmitC child pass execution failed.\n"; + return failure(); + } + return success(); +} + +static LogicalResult compileMixedEmitCChild(const BackendChildPlan &plan, + PTOBuildLevel effectiveLevel, + llvm::StringRef arch, + CompiledBackendChild &compiled, + llvm::raw_ostream &diagOS) { + ModuleOp module = plan.module; + if (failed(runEmitCBackendPipeline(module, effectiveLevel, arch, diagOS))) + return failure(); + + compiled.plan = &plan; + compiled.backend = PTOBackend::EmitC; + compiled.kernelKind = plan.kernelKind; + if (failed(emitEmitCSourceArtifact(module, emitAddPtrTrace, + compiled.cppSource, diagOS))) + return failure(); + return success(); +} + +static LogicalResult compileMixedVPTOChild(const BackendChildPlan &plan, + PTOBuildLevel effectiveLevel, + llvm::StringRef arch, int argc, + char **argv, + CompiledBackendChild &compiled, + llvm::raw_ostream &diagOS) { + ModuleOp module = plan.module; + const bool hasTileOpsToExpand = hasUnexpandedTileOps(module); + const bool hasTilelangHelpers = hasTilelangInlineHelpers(module); + + if (hasTileOpsToExpand) { + if (failed(runSharedPTOToSeamPipeline(module, effectiveLevel, arch, diagOS))) + return failure(); + } + + if (failed(runVPTOBackendPipeline(module, argc, argv, hasTileOpsToExpand, + hasTilelangHelpers))) + return failure(); + + compiled.plan = &plan; + compiled.backend = PTOBackend::VPTO; + compiled.kernelKind = plan.kernelKind; + compiled.vptoModule = module; + + pto::VPTOEmissionOptions options = buildVPTOEmissionOptions(); + if (failed(pto::lowerVPTOModuleToLLVMModules( + module, options, compiled.cubeModule, compiled.vectorModule, diagOS))) { + diagOS << "Error: Failed to lower mixed VPTO child to LLVM modules.\n"; + return failure(); + } + return success(); +} + +static LogicalResult compilePTOChildren(CompilationPlan &plan, + PTOBuildLevel effectiveLevel, + llvm::StringRef arch, int argc, + char **argv, + SmallVectorImpl + &compiledChildren, + llvm::raw_ostream &diagOS) { + for (BackendChildPlan &child : plan.children) { + if (failed(applyResolvedExternalImports(child))) + return failure(); + + CompiledBackendChild compiled; + if (child.backend == PTOBackend::EmitC) { + if (failed(compileMixedEmitCChild(child, effectiveLevel, arch, compiled, + diagOS))) + return failure(); + } else { + if (failed(compileMixedVPTOChild(child, effectiveLevel, arch, argc, argv, + compiled, diagOS))) + return failure(); + } + compiledChildren.push_back(std::move(compiled)); + } + return success(); +} + +static LogicalResult createTempPath(llvm::StringRef prefix, + llvm::StringRef suffix, + std::string &path, + llvm::raw_ostream &diagOS) { + llvm::SmallString<128> tempPath; + int fd = -1; + std::error_code ec = + llvm::sys::fs::createTemporaryFile(prefix, suffix, fd, tempPath); + if (ec) { + diagOS << "Error: failed to create temporary file for " << prefix << suffix + << ": " << ec.message() << "\n"; + return failure(); + } + llvm::sys::Process::SafelyCloseFileDescriptor(fd); + path = tempPath.str().str(); + return success(); +} + +static void removeTempPaths(ArrayRef paths) { + for (const std::string &path : paths) + if (!path.empty()) + llvm::sys::fs::remove(path); +} + +static LogicalResult compileMixedDeviceObjects( + ArrayRef compiledChildren, + const ObjectEmissionToolchain &toolchain, + SmallVectorImpl &artifacts, + SmallVectorImpl &tempPaths, llvm::raw_ostream &diagOS) { + for (const CompiledBackendChild &child : compiledChildren) { + if (child.backend == PTOBackend::EmitC) { + std::string cppPath; + std::string objPath; + std::string stderrPath; + if (failed(createTempPath("ptoas-mixed-emitc", ".cpp", cppPath, diagOS)) || + failed(createTempPath("ptoas-mixed-emitc", ".o", objPath, diagOS)) || + failed(createTempPath("ptoas-mixed-emitc", ".log", stderrPath, + diagOS))) + return failure(); + tempPaths.push_back(cppPath); + tempPaths.push_back(objPath); + tempPaths.push_back(stderrPath); + + if (failed(emitCppFatobj(child.cppSource, cppPath, objPath, toolchain, + stderrPath, diagOS))) + return failure(); + + BackendArtifact artifact; + artifact.backend = PTOBackend::EmitC; + artifact.kind = BackendArtifactKind::Fatobj; + artifact.path = objPath; + if (child.plan) + for (const BackendExportPlan &exportPlan : child.plan->exports) + artifact.exportedABISymbols.push_back(exportPlan.abiName); + artifacts.push_back(std::move(artifact)); + continue; + } + + auto emitVPTOObject = [&](llvm::Module *llvmModule, + ObjectEmissionDeviceTarget target, + BackendArtifactKind artifactKind) -> LogicalResult { + if (!llvmModule) + return success(); + std::string llPath; + std::string objPath; + std::string stderrPath; + if (failed(createTempPath("ptoas-mixed-vpto", ".ll", llPath, diagOS)) || + failed(createTempPath("ptoas-mixed-vpto", ".o", objPath, diagOS)) || + failed(createTempPath("ptoas-mixed-vpto", ".log", stderrPath, diagOS))) + return failure(); + tempPaths.push_back(llPath); + tempPaths.push_back(objPath); + tempPaths.push_back(stderrPath); + + if (failed(writeLLVMModule(*llvmModule, llPath, diagOS))) + return failure(); + if (failed(compileLLVMToDeviceObject(llPath, objPath, target, toolchain, + stderrPath, diagOS))) + return failure(); + + BackendArtifact artifact; + artifact.backend = PTOBackend::VPTO; + artifact.kind = artifactKind; + artifact.path = objPath; + if (child.plan) + for (const BackendExportPlan &exportPlan : child.plan->exports) + artifact.exportedABISymbols.push_back(exportPlan.abiName); + artifacts.push_back(std::move(artifact)); + return success(); + }; + + if (failed(emitVPTOObject(child.vectorModule.module.get(), + ObjectEmissionDeviceTarget::Vector, + BackendArtifactKind::VectorObject))) + return failure(); + if (failed(emitVPTOObject(child.cubeModule.module.get(), + ObjectEmissionDeviceTarget::Cube, + BackendArtifactKind::CubeObject))) + return failure(); + } + + if (artifacts.empty()) { + diagOS << "Error: mixed fatobj compilation produced no device objects.\n"; + return failure(); + } + return success(); +} + +static LogicalResult buildMixedHostStubSource( + ArrayRef compiledChildren, std::string &stubSource, + llvm::raw_ostream &diagOS) { + SmallVector stubModules; + for (const CompiledBackendChild &child : compiledChildren) { + if (child.backend != PTOBackend::VPTO || !child.vptoModule) + continue; + bool hasAICoreKernel = false; + ModuleOp vptoModule = child.vptoModule; + for (func::FuncOp func : vptoModule.getOps()) { + if (!func.isExternal() && func->hasAttr("pto.aicore")) { + hasAICoreKernel = true; + break; + } + } + if (hasAICoreKernel) + stubModules.push_back(child.vptoModule); + } + if (stubModules.empty()) { + stubSource = "#ifndef __global__\n#define __global__\n#endif\n\n" + "#ifndef __gm__\n#define __gm__\n#endif\n\n"; + return success(); + } + if (failed(pto::emitVPTOHostStubSource(stubModules, stubSource, diagOS))) { + diagOS << "Error: Failed to emit mixed VPTO host stub source.\n"; + return failure(); + } + return success(); +} + +static LogicalResult linkMixedFatobj( + ArrayRef artifacts, llvm::StringRef stubSource, + llvm::StringRef outputPath, const ObjectEmissionToolchain &toolchain, + SmallVectorImpl &tempPaths, llvm::raw_ostream &diagOS) { + std::string stubPath; + std::string stderrPath; + if (failed(createTempPath("ptoas-mixed-stub", ".cpp", stubPath, diagOS)) || + failed(createTempPath("ptoas-mixed", ".log", stderrPath, diagOS))) + return failure(); + tempPaths.push_back(stubPath); + tempPaths.push_back(stderrPath); + + if (failed(writeHostStubSource(stubSource, stubPath, diagOS))) + return failure(); + + std::string sanitizedModuleId = + outputPath.empty() || outputPath == "-" ? "ptoas_mixed_fatobj" + : outputPath.str(); + for (char &c : sanitizedModuleId) + if (!std::isalnum(static_cast(c)) && c != '_') + c = '_'; + + SmallVector fatobjPaths; + for (size_t i = 0, e = artifacts.size(); i < e; ++i) { + const BackendArtifact &artifact = artifacts[i]; + if (artifact.kind == BackendArtifactKind::Fatobj) { + fatobjPaths.push_back(artifact.path); + continue; + } + + std::string mergedObjPath; + std::string fatobjPath; + if (failed(createTempPath("ptoas-mixed-backend-merged", ".o", mergedObjPath, + diagOS)) || + failed(createTempPath("ptoas-mixed-backend-fatobj", ".o", fatobjPath, + diagOS))) + return failure(); + tempPaths.push_back(mergedObjPath); + tempPaths.push_back(fatobjPath); + + if (failed(mergeDeviceObjects(ArrayRef(artifact.path), + mergedObjPath, toolchain, + stderrPath, diagOS))) + return failure(); + + std::string groupModuleId = sanitizedModuleId + "_backend_" + + std::to_string(i); + if (failed(compileStubToFatobj(stubPath, mergedObjPath, fatobjPath, + groupModuleId, toolchain, stderrPath, + diagOS))) + return failure(); + fatobjPaths.push_back(fatobjPath); + } + + if (fatobjPaths.empty()) { + diagOS << "Error: mixed fatobj linking produced no backend fatobjs.\n"; + return failure(); + } + if (fatobjPaths.size() == 1) { + if (std::error_code ec = + llvm::sys::fs::copy_file(fatobjPaths.front(), outputPath)) { + diagOS << "Error: failed to copy mixed backend fatobj to " << outputPath + << ": " << ec.message() << "\n"; + return failure(); + } + return success(); + } + + if (failed(linkFatobjs(fatobjPaths, outputPath, toolchain, stderrPath, + diagOS))) + return failure(); + return success(); +} + +static int runMixedFatobjPlan(CompilationPlan &plan, + PTOBuildLevel effectiveLevel, + llvm::StringRef arch, int argc, char **argv, + llvm::StringRef outputPath, + llvm::raw_ostream &diagOS) { + if (emitMlirIR || emitVPTO || ptoPrintSeamIR || !ptoSeamIRFile.empty()) { + diagOS << "Error: mixed pto.backend fatobj mode does not support " + "debug IR output flags.\n"; + return 1; + } + if (outputPath.empty() || outputPath == "-") { + diagOS << "Error: mixed pto.backend fatobj mode requires an explicit " + "file path passed with -o.\n"; + return 1; + } + + SmallVector compiledChildren; + if (failed(compilePTOChildren(plan, effectiveLevel, arch, argc, argv, + compiledChildren, diagOS))) + return 1; + + std::string stubSource; + if (failed(buildMixedHostStubSource(compiledChildren, stubSource, diagOS))) + return 1; + + ObjectEmissionToolchain toolchain; + if (failed(discoverObjectEmissionToolchain(toolchain, diagOS))) + return 1; + + SmallVector tempPaths; + auto cleanup = llvm::make_scope_exit([&]() { removeTempPaths(tempPaths); }); + + SmallVector artifacts; + if (failed(compileMixedDeviceObjects(compiledChildren, toolchain, artifacts, + tempPaths, diagOS))) + return 1; + + if (failed(linkMixedFatobj(artifacts, stubSource, outputPath, toolchain, + tempPaths, diagOS))) + return 1; + return 0; +} + +static int emitVPTOBackendResult(ModuleOp module, + llvm::ToolOutputFile &outputFile) { + if (emitVPTO) { + module.print(outputFile.os()); + outputFile.os() << "\n"; + outputFile.keep(); + return 0; + } + + pto::VPTOEmissionOptions options = buildVPTOEmissionOptions(); + std::string stubSource; + if (failed(pto::emitVPTOHostStubSource(module, stubSource, llvm::errs()))) { + llvm::errs() << "Error: Failed to emit VPTO host stub source.\n"; + return 1; + } + + pto::EmittedLLVMModule cubeModule; + pto::EmittedLLVMModule vectorModule; + if (failed( + pto::lowerVPTOModuleToLLVMModules(module, options, cubeModule, + vectorModule, llvm::errs()))) { + llvm::errs() << "Error: Failed to lower VPTO to LLVM modules.\n"; + return 1; + } + + if (failed(pto::emitVPTOFatobjFromLLVMModules( + cubeModule.module.get(), vectorModule.module.get(), stubSource, + outputFile, llvm::errs()))) { + llvm::errs() << "Error: Failed to emit VPTO fatobj.\n"; + return 1; + } + outputFile.keep(); + return 0; +} + +static LogicalResult runVPTOBackendPipeline(ModuleOp module, int argc, + char **argv, + bool hasTileOpsToExpand, + bool hasTilelangHelpers) { + PassManager pm(module.getContext()); + pm.enableVerifier(); + pm.addPass(pto::createVPTOSplitCVModulePass()); + pm.addPass(pto::createVPTONormalizeContainerPass()); + if (!hasTileOpsToExpand && hasTilelangHelpers) + inlineTilelangHelpersOnVPTOInput(pm); + if (hasTileOpsToExpand) + lowerPTOToVPTOBackend(pm, argc, argv); + prepareVPTOForEmission(pm); + if (failed(applyConfiguredPassManagerCLOptions( + pm, "VPTO unified emission pipeline"))) + return failure(); + if (failed(pm.run(module))) { + llvm::errs() << "Error: VPTO emission pipeline failed.\n"; + return failure(); + } + if (failed(applyVPTOPublicABIExportNames(module))) + return failure(); + return success(); +} + +int mlir::pto::runPTOASDriver(int argc, char **argv) { + DialectRegistry registry; + registry.insert(); + registry.insert(); + registry.insert(); + registry.insert(); + registry.insert(); + registry.insert(); + registry.insert(); + registry.insert(); + registry.insert(); + + registry.insert(); + arith::registerBufferizableOpInterfaceExternalModels(registry); + tensor::registerBufferizableOpInterfaceExternalModels(registry); + pto::registerBufferizableOpInterfaceExternalModels(registry); + + registry.insert(); + registry.insert(); + mlir::registerAllPasses(); + ::registerPTOPasses(); + mlir::pto::registerPTOViewToMemrefPass(); + ::registerPTOInlineLibCall(); + ::registerFoldTileBufIntrinsics(); + ::registerExpandTileOp(); + mlir::registerPassManagerCLOptions(); + + llvm::cl::SetVersionPrinter(printPTOASVersion); + + bool cliArchSpecified = false; + for (int i = 1; i < argc; ++i) { + llvm::StringRef arg(argv[i]); + if (arg == "--pto-arch" || arg.starts_with("--pto-arch=")) + cliArchSpecified = true; + } + bool cliBackendSpecified = hasCLIOption(argc, argv, "--pto-backend"); + + // Parse command line options + llvm::cl::ParseCommandLineOptions(argc, argv, "PTO Assembler (ptoas)\n"); + + PTOBackend effectiveBackend = PTOBackend::EmitC; + if (!parseBackend(ptoBackend, effectiveBackend)) { + llvm::errs() << "Error: invalid --pto-backend='" << ptoBackend + << "'. Expected 'emitc' or 'vpto'.\n"; + return 1; + } + + // Read whole input first (so we can auto-detect .ptobc by magic). + auto fileOrErr = llvm::MemoryBuffer::getFileOrSTDIN(inputFilename); + if (!fileOrErr) { + llvm::errs() << "Error: Could not open input file: " + << fileOrErr.getError().message() << "\n"; + return 1; + } + + MLIRContext context(registry); + // Be tolerant: ptobc decode may materialize ops from dialects that aren't + // explicitly registered/loaded in this tool yet. + context.allowUnregisteredDialects(true); + + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.getOrLoadDialect(); + + OwningOpRef module; + llvm::StringRef buf = (*fileOrErr)->getBuffer(); + const bool isPTOBC = (buf.size() >= 6 && std::memcmp(buf.data(), "PTOBC\0", 6) == 0); + + auto normalizeArch = [](llvm::StringRef archValue) { + std::string normalized = archValue.str(); + for (char &c : normalized) + c = static_cast(std::tolower(static_cast(c))); + return normalized; + }; + auto detectTextualModuleArch = [&](llvm::StringRef text) -> std::optional { + llvm::SmallVector matches; + llvm::Regex archRegex( + R"ptoarch("?(pto\.target_arch)"?[[:space:]]*=[[:space:]]*"([[:alpha:][:digit:]_]+)")ptoarch"); + if (!archRegex.match(text, &matches) || matches.size() < 3) + return std::nullopt; + return normalizeArch(matches[2]); + }; + + std::string arch = normalizeArch(ptoTargetArch); + if (cliArchSpecified) { + if (arch != "a3" && arch != "a5") { + llvm::errs() << "Error: invalid --pto-arch='" << ptoTargetArch + << "'. Expected 'a3' or 'a5'.\n"; + return 1; + } + } else if (!isPTOBC) { + if (auto detectedArch = detectTextualModuleArch(buf)) + arch = *detectedArch; + } + if (arch != "a3" && arch != "a5") + arch = "a3"; + + if (isPTOBC) { + // Decode PTO bytecode directly into an MLIR module. + llvm::ArrayRef bytes(reinterpret_cast(buf.data()), buf.size()); +#if defined(__cpp_exceptions) || defined(__EXCEPTIONS) + try { + module = ptobc::decodePTOBCToModule(bytes, context); + } catch (...) { + llvm::errs() << "Error: Failed to decode PTOBC.\n"; + return 1; + } +#else + module = ptobc::decodePTOBCToModule(bytes, context); +#endif + if (!module) { + llvm::errs() << "Error: Failed to decode PTOBC.\n"; + return 1; + } + } else { + // Parse textual MLIR (.pto). + llvm::SourceMgr sourceMgr; + sourceMgr.AddNewSourceBuffer(std::move(*fileOrErr), llvm::SMLoc()); + pto::ScopedPTOParserTargetArch scopedParserArch( + &context, arch == "a5" ? pto::PTOParserTargetArch::A5 + : pto::PTOParserTargetArch::A3); + module = parseSourceFile(sourceMgr, &context); + if (!module) { + llvm::errs() << "Error: Failed to parse MLIR.\n"; + return 1; + } + } + + // If the CLI explicitly requested an arch, it overrides the input module. + // Otherwise, preserve the textual module's arch when present and only fall + // back to the effective default. + if (cliArchSpecified || !module->getOperation()->hasAttr("pto.target_arch")) { + module->getOperation()->setAttr("pto.target_arch", + mlir::StringAttr::get(&context, arch)); + } + + CompilationPlan compilationPlan; + if (failed(buildCompilationPlan(module.get(), cliBackendSpecified, + effectiveBackend, compilationPlan))) + return 1; + effectiveBackend = compilationPlan.effectiveBackend; + + if (effectiveBackend != PTOBackend::VPTO && + compilationPlan.mode != CompilationMode::MixedFatobj && + (emitVPTO || ptoPrintSeamIR || !ptoSeamIRFile.empty())) { + llvm::errs() << "Error: VPTO-specific flags require " + "--pto-backend=vpto or pto.backend = \"vpto\".\n"; + return 1; + } + + PTOBuildLevel effectiveLevel = defaultBuildLevel(); + if (!parseBuildLevel(ptoBuildLevel, effectiveLevel)) { + llvm::errs() << "Error: invalid --pto-level='" << ptoBuildLevel + << "'. Expected 'level1', 'level2', or 'level3'.\n"; + return 1; + } + + bool invalidAutoSyncTailHint = false; + module->walk([&](mlir::func::FuncOp func) { + auto hintAttr = + func->getAttrOfType("pto.auto_sync_tail_hint"); + if (!hintAttr) + return; + + std::string normalizedHint; + if (!parseAutoSyncTailHint(hintAttr.getValue(), normalizedHint)) { + func.emitError("invalid pto.auto_sync_tail_hint '") + << hintAttr.getValue() + << "'. Expected 'barrier-all' (or 'default') or " + "'mte3-to-s-event0'."; + invalidAutoSyncTailHint = true; + return; + } + func->setAttr("pto.auto_sync_tail_hint", + mlir::StringAttr::get(&context, normalizedHint)); + }); + if (invalidAutoSyncTailHint) + return 1; + + bool hasTAssign = false; + module->walk([&](pto::TAssignOp) { hasTAssign = true; }); + + if (hasTAssign && effectiveLevel != PTOBuildLevel::Level3) { + llvm::errs() << "Error: pto.tassign is only supported when " + "--pto-level=level3.\n"; + return 1; + } + + if (hasTAssign && enableInsertSync) { + llvm::errs() << "Error: pto.tassign requires --enable-insert-sync to be " + "disabled.\n"; + return 1; + } + + int enabledAutoSyncModes = + (enableInsertSync ? 1 : 0) + (enableInjectBarrierAllSync ? 1 : 0) + + (enableGraphSyncSolver ? 1 : 0); + if (enabledAutoSyncModes > 1) { + llvm::errs() << "Error: --enable-insert-sync, " + "--enable-inject-barrier-all-sync, and " + "--enable-graph-sync-solver are mutually exclusive.\n"; + return 1; + } + if (hasTAssign && enableInjectBarrierAllSync) { + llvm::errs() << "Error: pto.tassign requires " + "--enable-inject-barrier-all-sync to be disabled.\n"; + return 1; + } + if (hasTAssign && enableGraphSyncSolver) { + llvm::errs() << "Error: pto.tassign requires --enable-graph-sync-solver " + "to be disabled.\n"; + return 1; + } + + if (effectiveLevel == PTOBuildLevel::Level3) { + bool missing = false; + module->walk([&](pto::AllocTileOp op) { + if (!op.getAddr()) { + op.emitError("requires 'addr' operand when --pto-level=level3"); + missing = true; + } + }); + if (missing) + return 1; + } else { + bool hasAddr = false; + module->walk([&](pto::AllocTileOp op) { + if (op.getAddr()) { + op.emitError( + "unexpected 'addr' operand: only supported when --pto-level=level3"); + hasAddr = true; + } + }); + if (hasAddr) + return 1; + } + + if (compilationPlan.mode == CompilationMode::MixedFatobj) + return runMixedFatobjPlan(compilationPlan, effectiveLevel, arch, argc, argv, + outputFilename, llvm::errs()); + + // [Fix] ToolOutputFile Usage + std::error_code ec; + llvm::ToolOutputFile outputFile(outputFilename, ec, llvm::sys::fs::OF_None); + if (ec) { + llvm::errs() << ec.message() << "\n"; + return 1; + } + + const bool hasTileOpsToExpand = hasUnexpandedTileOps(*module); + const bool hasTilelangHelpers = hasTilelangInlineHelpers(*module); + + if (compilationPlan.mode == CompilationMode::VPTO && !hasTileOpsToExpand) { + if (ptoPrintSeamIR || !ptoSeamIRFile.empty()) { + llvm::errs() << "Error: shared pre-backend seam IR is unavailable when " + "skipping the shared PTO-to-VPTO lowering pipeline.\n"; + return 1; + } + if (failed(runVPTOBackendPipeline(module.get(), argc, argv, hasTileOpsToExpand, + hasTilelangHelpers))) + return 1; + return emitVPTOBackendResult(module.get(), outputFile); + } + + // Main PassManager + PassManager pm(&context); + + if (failed(configureSharedPTOPipeline(pm, effectiveLevel))) + return 1; + + llvm::raw_ostream *outputOS = &outputFile.os(); + + if (emitMlirIR) { + if (failed(pm.run(*module))) { + llvm::errs() << "Error: Pass execution failed.\n"; + return 1; + } + module->print(*outputOS); + outputFile.keep(); + return 0; + } + + addSharedPTOSeamPasses(pm); + if (failed(applyConfiguredPassManagerCLOptions(pm, "main PTOAS pipeline"))) + return 1; + + module->getOperation()->setAttr("pto.target_arch", + mlir::StringAttr::get(&context, arch)); + + if (compilationPlan.mode == CompilationMode::VPTO) { + if (failed(pm.run(*module))) { + llvm::errs() << "Error: Pass execution failed.\n"; + return 1; + } + + if (ptoPrintSeamIR) { + module->print(llvm::errs()); + llvm::errs() << "\n"; + } + if (failed(emitSharedPreBackendSeamIR(*module, ptoSeamIRFile))) + return 1; + + if (failed(runVPTOBackendPipeline(module.get(), argc, argv, hasTileOpsToExpand, + hasTilelangHelpers))) + return 1; + return emitVPTOBackendResult(module.get(), outputFile); + } + + if (arch == "a3") { + pm.addPass(pto::createEmitPTOManualPass(pto::PTOArch::A3)); + } else { + pm.addPass(pto::createEmitPTOManualPass(pto::PTOArch::A5)); + } + pm.addPass(emitc::createFormExpressionsPass()); + pm.addPass(mlir::createCSEPass()); + + if (failed(pm.run(*module))) { + llvm::errs() << "Error: Pass execution failed.\n"; + return 1; + } + + std::string cppOutput; + if (failed(emitEmitCSourceArtifact(module.get(), emitAddPtrTrace, cppOutput, + llvm::errs()))) + return 1; + + *outputOS << cppOutput; + outputOS->flush(); + + outputFile.keep(); // Success, keep the file + + return 0; +} diff --git a/tools/ptoas/VPTOFatobjEmission.h b/tools/ptoas/PTOASDriver.h similarity index 57% rename from tools/ptoas/VPTOFatobjEmission.h rename to tools/ptoas/PTOASDriver.h index fd161ba608..61dfbae99a 100644 --- a/tools/ptoas/VPTOFatobjEmission.h +++ b/tools/ptoas/PTOASDriver.h @@ -6,25 +6,12 @@ // 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. -#ifndef PTOAS_VPTO_FATOBJ_EMISSION_H -#define PTOAS_VPTO_FATOBJ_EMISSION_H - -#include "llvm/ADT/StringRef.h" -#include "mlir/Support/LogicalResult.h" - -namespace llvm { -class ToolOutputFile; -class Module; -class raw_ostream; -} +#ifndef PTOAS_DRIVER_H +#define PTOAS_DRIVER_H namespace mlir::pto { -LogicalResult emitVPTOFatobj(llvm::Module *cubeModule, - llvm::Module *vectorModule, - llvm::StringRef stubSource, - llvm::ToolOutputFile &outputFile, - llvm::raw_ostream &diagOS); +int runPTOASDriver(int argc, char **argv); } // namespace mlir::pto diff --git a/tools/ptoas/VPTOHostStubEmission.cpp b/tools/ptoas/VPTOHostStubEmission.cpp index ae56d85f8a..a2698ae3b3 100644 --- a/tools/ptoas/VPTOHostStubEmission.cpp +++ b/tools/ptoas/VPTOHostStubEmission.cpp @@ -74,35 +74,37 @@ static std::string getStubCType(Type type) { } // namespace static LogicalResult collectVPTOKernelStubDecls( - ModuleOp module, SmallVectorImpl &decls, + ArrayRef modules, SmallVectorImpl &decls, llvm::raw_ostream &diagOS) { bool hadError = false; llvm::StringMap logicalNameToIndex; - module.walk([&](func::FuncOp func) { - if (func.isExternal() || !hasVPTOKernelAttr(func)) - return; - - std::string logicalName = getLogicalKernelName(func.getSymName()); - SmallVector argTypes; - argTypes.reserve(func.getNumArguments()); - for (Type type : func.getArgumentTypes()) - argTypes.push_back(getStubCType(type)); - - auto [it, inserted] = - logicalNameToIndex.try_emplace(logicalName, decls.size()); - if (inserted) { - decls.push_back(VPTOKernelStubDecl{logicalName, std::move(argTypes)}); - return; - } - - VPTOKernelStubDecl &existing = decls[it->second]; - if (existing.argTypes != argTypes) { - diagOS << "Error: mixed kernel variants disagree on host stub signature " - << "for '" << logicalName << "'.\n"; - hadError = true; - } - }); + for (ModuleOp module : modules) { + module.walk([&](func::FuncOp func) { + if (func.isExternal() || !hasVPTOKernelAttr(func)) + return; + + std::string logicalName = getLogicalKernelName(func.getSymName()); + SmallVector argTypes; + argTypes.reserve(func.getNumArguments()); + for (Type type : func.getArgumentTypes()) + argTypes.push_back(getStubCType(type)); + + auto [it, inserted] = + logicalNameToIndex.try_emplace(logicalName, decls.size()); + if (inserted) { + decls.push_back(VPTOKernelStubDecl{logicalName, std::move(argTypes)}); + return; + } + + VPTOKernelStubDecl &existing = decls[it->second]; + if (existing.argTypes != argTypes) { + diagOS << "Error: mixed kernel variants disagree on host stub signature " + << "for '" << logicalName << "'.\n"; + hadError = true; + } + }); + } return hadError ? failure() : success(); } @@ -110,8 +112,14 @@ static LogicalResult collectVPTOKernelStubDecls( LogicalResult mlir::pto::emitVPTOHostStubSource(ModuleOp module, std::string &stubSource, llvm::raw_ostream &diagOS) { + return emitVPTOHostStubSource(ArrayRef(module), stubSource, diagOS); +} + +LogicalResult mlir::pto::emitVPTOHostStubSource(ArrayRef modules, + std::string &stubSource, + llvm::raw_ostream &diagOS) { SmallVector stubDecls; - if (failed(collectVPTOKernelStubDecls(module, stubDecls, diagOS))) + if (failed(collectVPTOKernelStubDecls(modules, stubDecls, diagOS))) return failure(); if (stubDecls.empty()) { diff --git a/tools/ptoas/VPTOHostStubEmission.h b/tools/ptoas/VPTOHostStubEmission.h index 7a091a503b..cf7c8fbfb6 100644 --- a/tools/ptoas/VPTOHostStubEmission.h +++ b/tools/ptoas/VPTOHostStubEmission.h @@ -24,6 +24,10 @@ namespace mlir::pto { LogicalResult emitVPTOHostStubSource(ModuleOp module, std::string &stubSource, llvm::raw_ostream &diagOS); +LogicalResult emitVPTOHostStubSource(ArrayRef modules, + std::string &stubSource, + llvm::raw_ostream &diagOS); + } // namespace mlir::pto #endif diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index caa32e5b36..b4b707c010 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -6,1677 +6,8 @@ // 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. -#include "PTO/IR/PTO.h" -#include "PTO/Transforms/VPTOLLVMEmitter.h" -#include "PTO/Transforms/Passes.h" -#include "PTO/Transforms/BufferizableOpInterfaceImpl.h" -#include "VPTOFatobjEmission.h" -#include "VPTOHostStubEmission.h" -#include "TilelangDaemon.h" -#include "mlir/IR/MLIRContext.h" -#include "mlir/IR/Diagnostics.h" -#include "mlir/IR/BuiltinOps.h" -#include "mlir/InitAllDialects.h" -#include "mlir/InitAllPasses.h" -#include "mlir/Parser/Parser.h" -#include "mlir/Pass/PassManager.h" -#include "mlir/Dialect/Affine/IR/AffineOps.h" -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/Dialect/Math/IR/Math.h" -#include -#include -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/Tensor/IR/Tensor.h" -#include "mlir/Target/Cpp/CppEmitter.h" -#include "llvm/Support/SourceMgr.h" -#include "llvm/Support/ToolOutputFile.h" -#include "llvm/Support/FileSystem.h" // [Fix] Required for OF_None -#include "llvm/Support/Path.h" -#include "ptobc/ptobc_decode.h" -#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h" -#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" -#include "mlir/Dialect/LLVMIR/LLVMDialect.h" -#include "mlir/Dialect/EmitC/IR/EmitC.h" -#include "mlir/Dialect/EmitC/Transforms/Passes.h" -#include "llvm/Support/CommandLine.h" -#include "llvm/Support/Regex.h" -#include "llvm/Support/raw_ostream.h" -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/SmallPtrSet.h" -#include "llvm/ADT/SmallVector.h" -#include "llvm/ADT/StringRef.h" -#include "llvm/ADT/StringSwitch.h" -#include "llvm/ADT/StringMap.h" -#include "llvm/Support/MemoryBuffer.h" -#include "llvm/Support/Program.h" -#include -#include -#include -#include -#include -#include -#include - -extern "C" { -extern char **environ; -} - -using namespace mlir; -using namespace pto; - -#ifndef PTOAS_RELEASE_VERSION -#define PTOAS_RELEASE_VERSION "unknown" -#endif - -int main(int argc, char **argv); - -static void printPTOASVersion(llvm::raw_ostream &os) { - os << "ptoas " << PTOAS_RELEASE_VERSION << "\n"; -} - -static std::string getParentDir(llvm::StringRef path) { - llvm::SmallString<256> parent(path); - llvm::sys::path::remove_filename(parent); - llvm::sys::path::remove_dots(parent, true); - return std::string(parent); -} - -static bool pathExists(llvm::StringRef path) { - return !path.empty() && llvm::sys::fs::exists(path); -} - -static std::string joinPath(llvm::StringRef lhs, llvm::StringRef rhs) { - llvm::SmallString<256> joined(lhs); - llvm::sys::path::append(joined, rhs); - llvm::sys::path::remove_dots(joined, true); - return std::string(joined); -} - -static std::string detectInstalledTilelangPath(const char *argv0) { - std::string exePath = llvm::sys::fs::getMainExecutable(argv0, (void *)&main); - if (exePath.empty()) - return {}; - - const std::string exeDir = getParentDir(exePath); - const std::string prefixDir = getParentDir(exeDir); - const std::string installedTileOps = joinPath(prefixDir, "share/ptoas/TileOps"); - if (pathExists(installedTileOps)) - return installedTileOps; - return {}; -} - -static std::string detectInstalledTilelangPkgPath(const char *argv0) { - std::string exePath = llvm::sys::fs::getMainExecutable(argv0, (void *)&main); - if (exePath.empty()) - return {}; - - const std::string exeDir = getParentDir(exePath); - const std::string prefixDir = getParentDir(exeDir); - const std::string installedPkgRoot = prefixDir; - const std::string installedPkg = joinPath(installedPkgRoot, "tilelang_dsl"); - if (pathExists(installedPkg)) - return installedPkgRoot; - return {}; -} - -static bool hasCLIOption(int argc, char **argv, llvm::StringRef option) { - const std::string optionWithValue = (option + "=").str(); - for (int i = 1; i < argc; ++i) { - llvm::StringRef arg(argv[i]); - if (arg == option || arg.starts_with(optionWithValue)) - return true; - } - return false; -} - -static LogicalResult applyConfiguredPassManagerCLOptions( - PassManager &pm, llvm::StringRef pipelineName, - llvm::raw_ostream &diagOS = llvm::errs()) { - if (succeeded(mlir::applyPassManagerCLOptions(pm))) - return success(); - diagOS << "Error: failed to apply MLIR pass manager command-line options for " - << pipelineName << ".\n"; - return failure(); -} - -static LogicalResult reorderEmitCFunctions(ModuleOp module) { - SmallVector declarations; - SmallVector definitions; - llvm::DenseMap definitionsByName; - - for (auto func : module.getOps()) { - if (func.isDeclaration()) { - declarations.push_back(func); - continue; - } - definitions.push_back(func); - definitionsByName[func.getSymNameAttr()] = func; - } - - llvm::DenseMap indegree; - llvm::DenseMap> outgoing; - for (auto func : definitions) - indegree[func.getOperation()] = 0; - - for (auto caller : definitions) { - Operation *callerOp = caller.getOperation(); - llvm::SmallPtrSet seenCallees; - bool hasCycle = false; - caller.walk([&](emitc::CallOp call) { - auto calleeAttr = call.getCalleeAttr(); - if (!calleeAttr) - return; - auto it = definitionsByName.find(calleeAttr.getLeafReference()); - if (it == definitionsByName.end()) - return; - Operation *calleeOp = it->second.getOperation(); - if (calleeOp == callerOp) { - hasCycle = true; - return; - } - if (!seenCallees.insert(calleeOp).second) - return; - outgoing[calleeOp].push_back(callerOp); - ++indegree[callerOp]; - }); - if (hasCycle) { - return caller.emitOpError() - << "recursive function calls are not supported for EmitC C++ " - "emission"; - } - } - - SmallVector ready; - for (auto func : definitions) { - if (indegree[func.getOperation()] == 0) - ready.push_back(func.getOperation()); - } - - SmallVector sortedDefinitions; - while (!ready.empty()) { - Operation *next = ready.front(); - ready.erase(ready.begin()); - auto nextFunc = cast(next); - sortedDefinitions.push_back(nextFunc); - - for (Operation *user : outgoing[next]) { - unsigned &userIndegree = indegree[user]; - if (--userIndegree == 0) - ready.push_back(user); - } - } - - if (sortedDefinitions.size() != definitions.size()) { - return module.emitError() - << "cyclic function call graph is not supported for EmitC C++ emission"; - } - - if (declarations.empty() && definitions.size() <= 1) - return success(); - - SmallVector desiredOrder; - desiredOrder.append(declarations.begin(), declarations.end()); - desiredOrder.append(sortedDefinitions.begin(), sortedDefinitions.end()); - - Block &body = module.getBodyRegion().front(); - Operation *anchor = nullptr; - for (Operation &op : body.getOperations()) { - if (isa(op)) { - anchor = &op; - break; - } - } - if (!anchor) - return success(); - - auto advanceAnchor = [&]() { - while (anchor) { - anchor = anchor->getNextNode(); - if (!anchor || isa(anchor)) - return; - } - }; - - for (auto func : desiredOrder) { - if (func.getOperation() == anchor) { - advanceAnchor(); - continue; - } - if (anchor) - func->moveBefore(anchor); - else - func->moveBefore(&body, body.end()); - } - - return success(); -} - -// -------------------------------------------------------------------------- -// Command Line Options -// -------------------------------------------------------------------------- -static llvm::cl::opt inputFilename(llvm::cl::Positional, - llvm::cl::desc(""), - llvm::cl::init("-")); - -static llvm::cl::opt outputFilename("o", - llvm::cl::desc("Output filename"), - llvm::cl::value_desc("filename"), - llvm::cl::init("-")); - -static llvm::cl::opt enableInsertSync("enable-insert-sync", - llvm::cl::desc("Enable automatic synchronization insertion pass"), - llvm::cl::init(false)); - -static llvm::cl::opt enableInjectBarrierAllSync( - "enable-inject-barrier-all-sync", - llvm::cl::desc("Enable conservative synchronization by inserting " - "pto.barrier PIPE_ALL before memory-effecting PTO pipe ops"), - llvm::cl::init(false)); - -static llvm::cl::opt enableGraphSyncSolver( - "enable-graph-sync-solver", - llvm::cl::desc("Enable the graph-based intra-core sync solver " - "(experimental). Mutually exclusive with " - "--enable-insert-sync and " - "--enable-inject-barrier-all-sync."), - llvm::cl::init(false)); - -static llvm::cl::opt graphSyncSolverEventIdMax( - "graph-sync-solver-event-id-max", - llvm::cl::desc( - "Maximum EVENT_ID slots for the graph sync solver (default 8). " - "Lower values exercise the PIPE_ALL coloring fallback sooner."), - llvm::cl::init(8)); - -static llvm::cl::opt enableTileOpExpand( - "enable-tile-op-expand", - llvm::cl::desc( - "Deprecated compatibility flag. TileOp expansion is controlled by " - "--pto-backend=vpto."), - llvm::cl::init(false)); - -#ifndef PTOAS_DEFAULT_TILELANG_PATH -#define PTOAS_DEFAULT_TILELANG_PATH "" -#endif -#ifndef PTOAS_DEFAULT_TILELANG_PKG_PATH -#define PTOAS_DEFAULT_TILELANG_PKG_PATH "" -#endif - -static llvm::cl::opt tilelangPath( - "tilelang-path", - llvm::cl::desc("Path to directory of .py tilelang DSL template files " - "(default: /lib/TileOps, baked in at build time)"), - llvm::cl::init(PTOAS_DEFAULT_TILELANG_PATH)); - -static llvm::cl::opt tilelangPkgPath( - "tilelang-pkg-path", - llvm::cl::desc("PYTHONPATH for tilelang_dsl package " - "(default: /tilelang-dsl/python, baked in at build time)"), - llvm::cl::init(PTOAS_DEFAULT_TILELANG_PKG_PATH)); - -static llvm::cl::opt daemonSocketPath( - "daemon-socket-path", - llvm::cl::desc("Path to Unix domain socket for daemon RPC " - "(default: /tmp/tilelang_daemon_{pid}.sock)"), - llvm::cl::init("")); - -static pto::ExpandTileOpOptions resolveExpandTileOpOptions(int argc, - char **argv) { - pto::ExpandTileOpOptions expandOpts; - expandOpts.tilelangPath = tilelangPath; - expandOpts.tilelangPkgPath = tilelangPkgPath; - - if (!hasCLIOption(argc, argv, "--tilelang-path")) { - std::string detectedTilelangPath = detectInstalledTilelangPath(argv[0]); - if (!detectedTilelangPath.empty()) - expandOpts.tilelangPath = detectedTilelangPath; - } - - if (!hasCLIOption(argc, argv, "--tilelang-pkg-path")) { - std::string detectedTilelangPkgPath = detectInstalledTilelangPkgPath(argv[0]); - if (!detectedTilelangPkgPath.empty()) - expandOpts.tilelangPkgPath = detectedTilelangPkgPath; - } - - // Daemon mode is default (no CLI option needed) - // Automatically start daemon for instance caching - if (!expandOpts.tilelangPath.empty()) { - std::string socket = daemonSocketPath; - if (socket.empty()) - socket = ptoas::DaemonManager::generateSocketPath(); - - // Register cleanup handler (daemon will be stopped on PTOAS exit) - ptoas::registerDaemonCleanup(); - - // Try to start daemon automatically - if (ptoas::DaemonManager::start(socket, expandOpts.tilelangPath, expandOpts.tilelangPkgPath)) { - expandOpts.daemonSocketPath = socket; - llvm::errs() << "Info: TileLang daemon started successfully\n"; - } else { - // Fallback: daemon failed, use subprocess mode (current approach) - expandOpts.daemonSocketPath = ""; - llvm::errs() << "Warning: Failed to start daemon, using subprocess mode (fallback)\n"; - } - } - - return expandOpts; -} - -static llvm::cl::opt disableInferLayout( - "disable-infer-layout", - llvm::cl::desc("Disable PTO layout inference pass (static-only)"), - llvm::cl::init(false)); - -static llvm::cl::opt emitAddPtrTrace( - "emit-addptr-trace", - llvm::cl::desc("Emit addptr trace comments in generated C++ output"), - llvm::cl::init(false)); - -static llvm::cl::opt emitMlirIR( - "emit-pto-ir", - llvm::cl::desc("Emit PTO IR after lowering instead of C++"), - llvm::cl::init(false)); - -static llvm::cl::opt ptoTargetArch( - "pto-arch", - llvm::cl::desc("Target Ascend architecture for codegen: a3 or a5 (default: a3)"), - llvm::cl::value_desc("a3|a5"), - llvm::cl::init("a3")); - -static llvm::cl::opt ptoBuildLevel( - "pto-level", - llvm::cl::desc("Build level for pass pipeline: level1, level2, or level3 (default: level2)"), - llvm::cl::value_desc("level1|level2|level3"), - llvm::cl::init("level2")); - -static llvm::cl::opt ptoBackend( - "pto-backend", - llvm::cl::desc("Final PTOAS backend: emitc or vpto (default: emitc)"), - llvm::cl::value_desc("emitc|vpto"), llvm::cl::init("emitc")); - -static llvm::cl::opt emitVPTO( - "emit-vpto", - llvm::cl::desc("Write final post-pass VPTO IR to -o"), - llvm::cl::init(false)); - -static llvm::cl::opt vptoPrintIR( - "vpto-print-ir", - llvm::cl::desc("Print post-pass VPTO backend IR to stderr"), - llvm::cl::init(false)); - -static llvm::cl::opt vptoLoweringStrategy( - "vpto-lowering-strategy", - llvm::cl::desc("VPTO vector lowering strategy: post-update or no-post-update"), - llvm::cl::value_desc("post-update|no-post-update"), - llvm::cl::init("post-update")); - -static llvm::cl::opt dumpVPTOIR( - "dump-vpto-ir", - llvm::cl::desc("Print post-pass VPTO backend IR to stderr"), - llvm::cl::init(false)); - -static llvm::cl::opt ptoPrintSeamIR( - "pto-print-seam-ir", - llvm::cl::desc("Print shared pre-backend seam IR to stderr"), - llvm::cl::init(false)); - -static llvm::cl::opt ptoSeamIRFile( - "pto-seam-ir-file", - llvm::cl::desc("Write shared pre-backend seam IR to a file"), - llvm::cl::value_desc("path"), - llvm::cl::init("")); - -enum class PTOBuildLevel { - Level1, - Level2, - Level3, -}; - -enum class PTOBackend { - EmitC, - VPTO, -}; - -static PTOBuildLevel defaultBuildLevel() { - return PTOBuildLevel::Level2; -} - -static bool parseBuildLevel(llvm::StringRef levelStr, PTOBuildLevel &out) { - std::string s = levelStr.str(); - for (char &c : s) - c = static_cast(std::tolower(static_cast(c))); - if (s == "level1") { - out = PTOBuildLevel::Level1; - return true; - } - if (s == "level2") { - out = PTOBuildLevel::Level2; - return true; - } - if (s == "level3") { - out = PTOBuildLevel::Level3; - return true; - } - return false; -} - -static constexpr llvm::StringLiteral kAutoSyncTailPolicyBarrierAll = - "barrier_all"; -static constexpr llvm::StringLiteral kAutoSyncTailPolicyMte3ToSEvent0 = - "setwait_mte3_to_s_event0"; - -static bool parseAutoSyncTailHint(llvm::StringRef hintStr, std::string &normalized) { - std::string s = hintStr.str(); - for (char &c : s) - c = static_cast(std::tolower(static_cast(c))); - if (s == "barrier-all" || s == "barrier_all" || s == "default") { - normalized = kAutoSyncTailPolicyBarrierAll.str(); - return true; - } - if (s == "mte3-to-s-event0" || s == "mte3_to_s_event0" || - s == "setwait-mte3-to-s-event0" || - s == "setwait_mte3_to_s_event0") { - normalized = kAutoSyncTailPolicyMte3ToSEvent0.str(); - return true; - } - return false; -} - -static bool parseBackend(llvm::StringRef backendStr, PTOBackend &out) { - std::string s = backendStr.str(); - for (char &c : s) - c = static_cast(std::tolower(static_cast(c))); - if (s == "emitc") { - out = PTOBackend::EmitC; - return true; - } - if (s == "vpto") { - out = PTOBackend::VPTO; - return true; - } - return false; -} - -static LogicalResult emitSharedPreBackendSeamIR(ModuleOp module, - llvm::StringRef outputPath) { - if (outputPath.empty()) - return success(); - - if (outputPath == "-") { - module->print(llvm::outs()); - llvm::outs() << "\n"; - llvm::outs().flush(); - return success(); - } - - std::error_code ec; - llvm::ToolOutputFile outputFile(outputPath, ec, llvm::sys::fs::OF_None); - if (ec) { - llvm::errs() << "Error: failed to open seam IR file '" << outputPath - << "': " << ec.message() << "\n"; - return failure(); - } - - module->print(outputFile.os()); - outputFile.os() << "\n"; - outputFile.keep(); - return success(); -} - -static bool hasUnexpandedTileOps(ModuleOp module) { - bool found = false; - module.walk([&](Operation *op) { - if (found) - return; - if (isa(op)) - found = true; - }); - return found; -} - -static bool hasTilelangInlineHelpers(ModuleOp module) { - bool found = false; - module.walk([&](func::FuncOp func) { - if (found) - return; - if (func->hasAttr("pto.tilelang.inline_proc")) - found = true; - }); - return found; -} - -// -------------------------------------------------------------------------- -// Post-process C++ output: rewrite marker calls into Tile member calls. -// -// We emit marker calls in EmitC IR because EmitC currently does not provide a -// first-class op for member-function invocation. After translation, we rewrite: -// PTOAS__TILE_SET_VALUE(dst, offset, val) -> dst.SetValue(offset, val) -// PTOAS__TILE_GET_VALUE(src, offset) -> src.GetValue(offset) -// PTOAS__TILE_DATA(obj) -> obj.data() -// PTOAS__TILE_SET_VALIDSHAPE(obj, r, c) -> obj.SetValidShape(r, c) -// PTOAS__PTR_LOAD(ptr, offset) -> ptr[offset] -// PTOAS__PTR_STORE(ptr, offset, val) -> ptr[offset] = val -// PTOAS__EVENTID_ARRAY_LOAD(arr, idx) -> arr[idx] -// PTOAS__EVENTID_ARRAY_STORE(arr, idx, v) -> arr[idx] = v -// -------------------------------------------------------------------------- -struct ParsedMarkerCall { - size_t markerPos = std::string::npos; - size_t rparenPos = std::string::npos; - llvm::SmallVector args; -}; - -struct MarkerRewriteSpec { - llvm::StringRef marker; - llvm::StringRef memberName; - unsigned expectedNumArgs = 0; -}; - -struct MarkerSubscriptRewriteSpec { - llvm::StringRef marker; - unsigned expectedNumArgs = 0; - bool isStore = false; -}; - -static bool parseMarkerArgs(llvm::StringRef argsRef, - llvm::SmallVectorImpl &args) { - size_t partBegin = 0; - int parenDepth = 0; - for (size_t i = 0; i < argsRef.size(); ++i) { - char c = argsRef[i]; - if (c == '(') { - ++parenDepth; - continue; - } - if (c == ')') { - if (parenDepth > 0) - --parenDepth; - continue; - } - if (c == ',' && parenDepth == 0) { - args.push_back(argsRef.slice(partBegin, i).trim()); - partBegin = i + 1; - } - } - if (partBegin > argsRef.size()) - return false; - args.push_back(argsRef.drop_front(partBegin).trim()); - return true; -} - -static std::optional -findNextMarkerCall(const std::string &cpp, llvm::StringRef marker, - size_t searchPos) { - ParsedMarkerCall call; - call.markerPos = cpp.find(marker.str(), searchPos); - if (call.markerPos == std::string::npos) - return std::nullopt; - - size_t lparenPos = call.markerPos + marker.size(); - if (lparenPos >= cpp.size() || cpp[lparenPos] != '(') - return ParsedMarkerCall{call.markerPos, std::string::npos, {}}; - - size_t argsBegin = lparenPos + 1; - int parenDepth = 0; - for (size_t i = argsBegin; i < cpp.size(); ++i) { - char c = cpp[i]; - if (c == '(') { - ++parenDepth; - continue; - } - if (c != ')') - continue; - if (parenDepth == 0) { - call.rparenPos = i; - break; - } - --parenDepth; - } - if (call.rparenPos == std::string::npos) - return call; - - llvm::StringRef argsRef(cpp.data() + argsBegin, call.rparenPos - argsBegin); - if (!parseMarkerArgs(argsRef, call.args)) - call.args.clear(); - return call; -} - -template -static bool rewriteMarkerCalls(std::string &cpp, llvm::StringRef marker, - BuildReplacementFn buildReplacement) { - size_t searchPos = 0; - bool changed = false; - for (auto call = findNextMarkerCall(cpp, marker, searchPos); call; - call = findNextMarkerCall(cpp, marker, searchPos)) { - if (call->rparenPos == std::string::npos) { - searchPos = call->markerPos + marker.size(); - continue; - } - - std::optional replacement = buildReplacement(*call); - if (!replacement) { - searchPos = call->rparenPos + 1; - continue; - } - - cpp.replace(call->markerPos, (call->rparenPos - call->markerPos) + 1, - *replacement); - changed = true; - searchPos = call->markerPos + replacement->size(); - } - return changed; -} - -static bool rewriteMarkerCallToMember(std::string &cpp, llvm::StringRef marker, - llvm::StringRef memberName, - unsigned expectedNumArgs) { - return rewriteMarkerCalls( - cpp, marker, [&](const ParsedMarkerCall &call) -> std::optional { - if (call.args.size() != expectedNumArgs) - return std::nullopt; - - std::string replacement; - replacement.reserve(marker.size() + 16); - replacement.append(call.args[0].str()); - replacement.push_back('.'); - replacement.append(memberName.str()); - replacement.push_back('('); - if (expectedNumArgs >= 2) - replacement.append(call.args[1].str()); - if (expectedNumArgs == 3) { - replacement.append(", "); - replacement.append(call.args[2].str()); - } - replacement.push_back(')'); - return replacement; - }); -} - -static void rewriteMarkerCallsToMembers( - std::string &cpp, llvm::ArrayRef rewrites) { - bool changed = true; - while (changed) { - changed = false; - for (const MarkerRewriteSpec &rewrite : rewrites) { - changed |= rewriteMarkerCallToMember(cpp, rewrite.marker, - rewrite.memberName, - rewrite.expectedNumArgs); - } - } -} - -static void rewriteTileGetSetValueMarkers(std::string &cpp) { - static const MarkerRewriteSpec kTileMarkerRewrites[] = { - {"PTOAS__TILE_SET_VALUE", "SetValue", 3}, - {"PTOAS__TILE_GET_VALUE", "GetValue", 2}, - {"PTOAS__TILE_DATA", "data", 1}, - {"PTOAS__TILE_SET_VALIDSHAPE", "SetValidShape", 3}, - }; - rewriteMarkerCallsToMembers(cpp, kTileMarkerRewrites); -} - -static void rewriteAsyncEventMarkers(std::string &cpp) { - static const MarkerRewriteSpec kAsyncEventMarkerRewrites[] = { - {"PTOAS__ASYNC_EVENT_WAIT", "Wait", 2}, - {"PTOAS__ASYNC_EVENT_TEST", "Test", 2}, - }; - rewriteMarkerCallsToMembers(cpp, kAsyncEventMarkerRewrites); -} - -// -------------------------------------------------------------------------- -// EmitC cleanup: drop empty emitc.expression ops. -// -// After FormExpressions + CSE, EmitC expressions can become empty when their -// root op is CSE'd with an equivalent dominating value outside the expression -// region. Such expressions crash mlir::emitc::translateToCpp because -// ExpressionOp::getRootOp() returns nullptr. -// -------------------------------------------------------------------------- -static void dropEmptyEmitCExpressions(Operation *rootOp) { - llvm::SmallVector toErase; - rootOp->walk([&](emitc::ExpressionOp expr) { - if (expr.getRootOp()) - return; - Block *body = expr.getBody(); - if (!body) - return; - auto yield = dyn_cast(body->getTerminator()); - if (!yield || yield.getNumOperands() != 1) - return; - Value yielded = yield.getOperand(0); - expr.getResult().replaceAllUsesWith(yielded); - toErase.push_back(expr); - }); - for (emitc::ExpressionOp expr : llvm::reverse(toErase)) - expr.erase(); -} - -static Attribute getDefaultEmitCVariableInitAttr(OpBuilder &builder, Type type) { - if (auto intTy = dyn_cast(type)) - return builder.getIntegerAttr(intTy, 0); - if (isa(type)) - return builder.getIndexAttr(0); - if (auto floatTy = dyn_cast(type)) - return builder.getFloatAttr(floatTy, 0.0); - if (isa(type)) - return emitc::OpaqueAttr::get(builder.getContext(), ""); - return Attribute{}; -} - -// FormExpressions may inline conditions into emitc.expression, but the C++ -// emitter prints cf.br/cf.cond_br operands by variable name rather than by -// recursively emitting an expression. Materialize such operands so CFG-based -// lowering (e.g. scf.while -> cf.*) stays valid. -static void materializeControlFlowOperands(Operation *rootOp) { - llvm::SmallVector branches; - rootOp->walk([&](Operation *op) { - if (isa(op)) - branches.push_back(op); - }); - - OpBuilder builder(rootOp->getContext()); - for (Operation *op : branches) { - builder.setInsertionPoint(op); - for (OpOperand &operand : op->getOpOperands()) { - Value value = operand.get(); - auto expr = dyn_cast_or_null(value.getDefiningOp()); - if (!expr) - continue; - - Attribute initAttr = - getDefaultEmitCVariableInitAttr(builder, value.getType()); - if (!initAttr) - continue; - - Value tmp = - builder.create(op->getLoc(), value.getType(), - initAttr) - .getResult(); - builder.create(op->getLoc(), tmp, value); - operand.set(tmp); - } - } -} - -static bool rewriteMarkerCallToSubscript(std::string &cpp, llvm::StringRef marker, - unsigned expectedNumArgs, - bool isStore) { - return rewriteMarkerCalls( - cpp, marker, [&](const ParsedMarkerCall &call) -> std::optional { - if (call.args.size() != expectedNumArgs) - return std::nullopt; - if (isStore) { - return (call.args[0] + "[" + call.args[1] + "] = " + call.args[2]) - .str(); - } - return (call.args[0] + "[" + call.args[1] + "]").str(); - }); -} - -static void rewriteMarkerCallsToSubscripts( - std::string &cpp, llvm::ArrayRef rewrites) { - bool changed = true; - while (changed) { - changed = false; - for (const MarkerSubscriptRewriteSpec &rewrite : rewrites) { - changed |= rewriteMarkerCallToSubscript(cpp, rewrite.marker, - rewrite.expectedNumArgs, - rewrite.isStore); - } - } -} - -static void rewritePtrScalarMarkers(std::string &cpp) { - static const MarkerSubscriptRewriteSpec kPtrMarkerRewrites[] = { - {"PTOAS__PTR_LOAD", 2, false}, - {"PTOAS__PTR_STORE", 3, true}, - }; - rewriteMarkerCallsToSubscripts(cpp, kPtrMarkerRewrites); -} - -static void rewriteEventIdArrayMarkers(std::string &cpp) { - static const MarkerSubscriptRewriteSpec kEventIdMarkerRewrites[] = { - {"PTOAS__EVENTID_ARRAY_LOAD", 2, false}, - {"PTOAS__EVENTID_ARRAY_STORE", 3, true}, - }; - rewriteMarkerCallsToSubscripts(cpp, kEventIdMarkerRewrites); -} - -static bool rewriteAddPtrTraceMarkers(std::string &cpp, bool showTrace) { - size_t searchPos = 0; - bool changed = false; - for (auto call = findNextMarkerCall(cpp, "PTOAS__ADDPTR_TRACE", searchPos); - call; call = findNextMarkerCall(cpp, "PTOAS__ADDPTR_TRACE", searchPos)) { - if (call->rparenPos == std::string::npos) { - searchPos = call->markerPos + 1; - continue; - } - if (call->args.size() != 3) { - searchPos = call->rparenPos + 1; - continue; - } - - std::string replacement; - if (showTrace) { - replacement.reserve(64); - replacement.append("/* ADDPTR_TRACE: "); - replacement.append(call->args[0].str()); - replacement.append(" = "); - replacement.append(call->args[1].str()); - replacement.append(" + "); - replacement.append(call->args[2].str()); - replacement.append(" */"); - } - - size_t replaceEnd = call->rparenPos; - if (!showTrace) { - size_t i = call->rparenPos + 1; - while (i < cpp.size() && std::isspace(static_cast(cpp[i]))) - ++i; - if (i < cpp.size() && cpp[i] == ';') - replaceEnd = i; - } - - cpp.replace(call->markerPos, (replaceEnd - call->markerPos) + 1, - replacement); - changed = true; - searchPos = call->markerPos + replacement.size(); - } - return changed; -} - -static bool isGeneratedGlobalTensorDecl(llvm::StringRef trimmed, - llvm::StringRef &decl, - llvm::StringRef &varName) { - if (!trimmed.starts_with("GlobalTensor<") || !trimmed.ends_with(";") || - trimmed.contains('=') || trimmed.contains('(')) { - return false; - } - - decl = trimmed.drop_back().rtrim(); - size_t lastWs = decl.find_last_of(" \t"); - if (lastWs == llvm::StringRef::npos) - return false; - varName = decl.drop_front(lastWs + 1); - if (!varName.starts_with("v") || varName.size() <= 1) - return false; - return llvm::all_of(varName.drop_front(1), - [](char c) { return std::isdigit(c); }); -} - -static void rewriteHoistedGlobalTensorDecls(std::string &cpp) { - // When `declareVariablesAtTop` is enabled, the C++ emitter hoists SSA value - // declarations to the top of the function and emits assignments later. This - // requires the C++ type to be default-constructible. - // - // `GlobalTensor<...>` from pto-isa does NOT have a default constructor, so a - // hoisted declaration like: - // GlobalTensor<...> v42; - // fails to compile. Initialize those hoisted temporaries with a null pointer - // so they are constructible: - // GlobalTensor<...> v42(nullptr); - // - // We keep the assignment later; the null-initialized value is never used. - std::string out; - out.reserve(cpp.size() + 64); - - llvm::StringRef ref(cpp); - while (!ref.empty()) { - auto split = ref.split('\n'); - llvm::StringRef line = split.first; - llvm::StringRef rest = split.second; - - llvm::StringRef trimmed = line.trim(); - bool rewritten = false; - llvm::StringRef decl; - llvm::StringRef varName; - if (isGeneratedGlobalTensorDecl(trimmed, decl, varName)) { - size_t indentLen = line.find_first_not_of(" \t"); - if (indentLen == std::string::npos) - indentLen = 0; - llvm::StringRef indent = line.take_front(indentLen); - - out.append(indent.str()); - out.append(decl.str()); - out.append("(nullptr);"); - rewritten = true; - } - - if (!rewritten) - out.append(line.str()); - if (!rest.empty()) - out.push_back('\n'); - ref = rest; - } - - cpp.swap(out); -} - -namespace { -struct ConstantDeclCandidate { - size_t declLine = 0; - std::string indent; - std::string type; - bool hasInitializer = false; - std::string initializer; - size_t assignmentCount = 0; - size_t assignmentLine = 0; - std::string assignmentRhs; -}; -} // namespace - -static bool isGeneratedValueName(llvm::StringRef name) { - if (!name.consume_front("v") || name.empty()) - return false; - return llvm::all_of(name, [](char c) { return std::isdigit(c); }); -} - -static bool isConstFoldableScalarType(llvm::StringRef type) { - type = type.trim(); - if (type.starts_with("const ") || type.starts_with("constexpr ")) - return false; - return llvm::StringSwitch(type) - .Cases("bool", "float", "double", "half", "bfloat16_t", true) - .Cases("int8_t", "uint8_t", "int16_t", "uint16_t", true) - .Cases("int32_t", "uint32_t", "int64_t", "uint64_t", true) - .Default(false); -} - -static bool isLiteralInitializer(llvm::StringRef rhs) { - rhs = rhs.trim(); - if (rhs.empty()) - return false; - if (rhs == "true" || rhs == "false" || rhs == "nullptr") - return true; - - static const llvm::Regex kIntLiteral( - R"(^[+-]?(0[xX][0-9A-Fa-f]+|[0-9]+)[uUlL]*$)"); - static const llvm::Regex kFloatLiteral( - R"(^[+-]?(([0-9]+\.[0-9]*|\.[0-9]+|[0-9]+)([eE][+-]?[0-9]+)?|[0-9]+[eE][+-]?[0-9]+)[fF]?$)"); - static const llvm::Regex kHexFloatLiteral( - R"(^[+-]?0[xX]([0-9A-Fa-f]+\.[0-9A-Fa-f]*|[0-9A-Fa-f]+|\.[0-9A-Fa-f]+)[pP][+-]?[0-9]+[fF]?$)"); - static const llvm::Regex kSpecialFloatLiteral( - R"(^[+-]?(nan|inf)[fF]?$)"); - - return kIntLiteral.match(rhs) || kFloatLiteral.match(rhs) || - kHexFloatLiteral.match(rhs) || kSpecialFloatLiteral.match(rhs); -} - -static std::string normalizeConstInitializer(llvm::StringRef type, - llvm::StringRef rhs) { - type = type.trim(); - rhs = rhs.trim(); - if (type == "bool") { - if (rhs == "0" || rhs == "false") - return "false"; - if (rhs == "1" || rhs == "-1" || rhs == "true") - return "true"; - } - return rhs.str(); -} - -static bool parseConstantDeclarationLine(llvm::StringRef line, - ConstantDeclCandidate &candidate, - std::string &valueName) { - llvm::StringRef trimmed = line.trim(); - if (trimmed.empty() || trimmed.starts_with("#") || trimmed.starts_with("//") || - !trimmed.ends_with(";")) - return false; - - llvm::StringRef body = trimmed.drop_back().rtrim(); - if (body.starts_with("return") || body.starts_with("goto ") || - body.starts_with("if ") || body.starts_with("if(") || - body.starts_with("switch ") || body.starts_with("switch(") || - body.starts_with("for ") || body.starts_with("for(") || - body.starts_with("while ") || body.starts_with("while(") || - body.starts_with("case ") || body == "default") - return false; - - llvm::StringRef lhs = body; - llvm::StringRef rhs; - if (size_t eqPos = body.find('='); eqPos != llvm::StringRef::npos) { - lhs = body.take_front(eqPos).rtrim(); - rhs = body.drop_front(eqPos + 1).trim(); - } - - size_t lastWs = lhs.find_last_of(" \t"); - if (lastWs == llvm::StringRef::npos) - return false; - - llvm::StringRef type = lhs.take_front(lastWs).rtrim(); - llvm::StringRef name = lhs.drop_front(lastWs + 1).trim(); - if (!isGeneratedValueName(name) || !isConstFoldableScalarType(type)) - return false; - - size_t indentLen = line.find_first_not_of(" \t"); - if (indentLen == llvm::StringRef::npos) - indentLen = 0; - candidate.indent = line.take_front(indentLen).str(); - candidate.type = type.str(); - valueName = name.str(); - - if (!rhs.empty()) { - if (!isLiteralInitializer(rhs)) - return false; - candidate.hasInitializer = true; - candidate.initializer = normalizeConstInitializer(type, rhs); - } - - return true; -} - -static bool parseGeneratedValueAssignment(llvm::StringRef line, - llvm::StringRef &valueName, - llvm::StringRef &rhs) { - llvm::StringRef trimmed = line.trim(); - if (trimmed.empty() || trimmed.starts_with("#") || trimmed.starts_with("//") || - !trimmed.ends_with(";")) - return false; - - llvm::StringRef body = trimmed.drop_back().rtrim(); - size_t eqPos = body.find('='); - if (eqPos == llvm::StringRef::npos) - return false; - - llvm::StringRef lhs = body.take_front(eqPos).rtrim(); - rhs = body.drop_front(eqPos + 1).trim(); - if (!isGeneratedValueName(lhs)) - return false; - valueName = lhs; - return true; -} - -static void rewriteScalarConstantDecls(std::string &cpp) { - llvm::SmallVector lines; - for (llvm::StringRef ref(cpp); !ref.empty(); ref = ref.split('\n').second) { - auto split = ref.split('\n'); - lines.push_back(split.first.str()); - } - - llvm::SmallVector eraseLine(lines.size(), false); - auto rewriteSegment = [&](size_t beginLine, size_t endLine) { - llvm::StringMap candidates; - - for (size_t i = beginLine; i <= endLine; ++i) { - ConstantDeclCandidate candidate; - std::string valueName; - if (parseConstantDeclarationLine(lines[i], candidate, valueName)) { - candidate.declLine = i; - candidates[valueName] = std::move(candidate); - continue; - } - - llvm::StringRef assignedName; - llvm::StringRef rhs; - if (!parseGeneratedValueAssignment(lines[i], assignedName, rhs)) - continue; - - auto it = candidates.find(assignedName); - if (it == candidates.end()) - continue; - - ConstantDeclCandidate &info = it->second; - ++info.assignmentCount; - info.assignmentLine = i; - info.assignmentRhs = rhs.str(); - } - - for (auto &entry : candidates) { - llvm::StringRef valueName = entry.getKey(); - ConstantDeclCandidate &info = entry.getValue(); - - std::string initializer; - if (info.hasInitializer) { - if (info.assignmentCount != 0) - continue; - initializer = info.initializer; - } else { - if (info.assignmentCount != 1) - continue; - if (!isLiteralInitializer(info.assignmentRhs)) - continue; - initializer = normalizeConstInitializer( - info.type, llvm::StringRef(info.assignmentRhs)); - eraseLine[info.assignmentLine] = true; - } - - lines[info.declLine] = (info.indent + "const " + info.type + " " + - valueName.str() + " = " + initializer + ";"); - } - }; - - int braceDepth = 0; - size_t segmentStart = 0; - for (size_t i = 0; i < lines.size(); ++i) { - int depthBefore = braceDepth; - for (char c : lines[i]) { - if (c == '{') - ++braceDepth; - else if (c == '}') - --braceDepth; - } - - if (depthBefore == 0 && braceDepth > 0) - segmentStart = i; - if (depthBefore > 0 && braceDepth == 0) - rewriteSegment(segmentStart, i); - } - - std::string out; - out.reserve(cpp.size()); - for (size_t i = 0; i < lines.size(); ++i) { - if (eraseLine[i]) - continue; - out.append(lines[i]); - if (i + 1 != lines.size()) - out.push_back('\n'); - } - cpp.swap(out); -} - -static bool shouldDeclareVariablesAtTop(ModuleOp module) { - auto hasMultiBlockFunc = [](auto func) { return func.getBlocks().size() > 1; }; - return llvm::any_of(module.getOps(), hasMultiBlockFunc) || - llvm::any_of(module.getOps(), hasMultiBlockFunc); -} - -static void prepareVPTOForEmission(PassManager &pm) { - auto &kernelModulePM = pm.nest(); - kernelModulePM.addPass(createCanonicalizerPass()); - kernelModulePM.addPass(createCSEPass()); - kernelModulePM.addPass(pto::createVPTOPtrNormalizePass()); - kernelModulePM.addPass(pto::createVPTOPtrCastCleanupPass()); - kernelModulePM.addPass(createReconcileUnrealizedCastsPass()); - kernelModulePM.addNestedPass( - createVPTOExpandWrapperOpsPass()); - kernelModulePM.addPass(createCSEPass()); - kernelModulePM.addNestedPass( - pto::createPTOInferVPTOVecScopePass()); - kernelModulePM.addPass(createCanonicalizerPass()); - kernelModulePM.addPass(createCSEPass()); - kernelModulePM.addPass(pto::createPTOValidateVPTOEmissionIRPass()); -} - -static void lowerPTOToVPTOBackend(PassManager &pm, int argc, char **argv) { - // TileOp Expand path: - // 1. ExpandTileOp: instantiate TileLang DSL templates, replace tile ops - // with func.call to template functions (tile_buf params) - // 2. InlineLibCall: inline template function bodies - // 3. FoldTileBufIntrinsics: fold tile_buf_addr / tile_valid_rows / - // tile_valid_cols to concrete memref/constant values - auto &kernelModulePM = pm.nest(); - pto::ExpandTileOpOptions expandOpts = resolveExpandTileOpOptions(argc, argv); - kernelModulePM.addPass(pto::createExpandTileOpPass(expandOpts)); - - kernelModulePM.addPass(pto::createPTOInlineLibCallPass()); - kernelModulePM.addNestedPass( - pto::createFoldTileBufIntrinsicsPass()); - // FoldTileBufIntrinsics materializes many constant branch conditions. - // Clean them up immediately on the TileOp expansion path before the - // authoring-stage VPTO verifier and let the existing CSE passes remove the - // resulting dead values later in the pipeline. - kernelModulePM.addPass(mlir::createSCCPPass()); - kernelModulePM.addPass(mlir::createCanonicalizerPass()); -} - -static void inlineTilelangHelpersOnVPTOInput(PassManager &pm) { - auto &kernelModulePM = pm.nest(); - kernelModulePM.addPass(pto::createPTOInlineLibCallPass()); - kernelModulePM.addPass(mlir::createSCCPPass()); - kernelModulePM.addPass(mlir::createCanonicalizerPass()); -} - -static pto::VPTOEmissionOptions buildVPTOEmissionOptions() { - pto::VPTOEmissionOptions options; - options.dumpVPTOIR = false; - options.targetTriple = "hiipu64-hisilicon-cce"; - return options; -} - -static int emitVPTOBackendResult(ModuleOp module, - llvm::ToolOutputFile &outputFile) { - if (emitVPTO) { - module.print(outputFile.os()); - outputFile.os() << "\n"; - outputFile.keep(); - return 0; - } - - pto::VPTOEmissionOptions options = buildVPTOEmissionOptions(); - std::string stubSource; - if (failed(pto::emitVPTOHostStubSource(module, stubSource, llvm::errs()))) { - llvm::errs() << "Error: Failed to emit VPTO host stub source.\n"; - return 1; - } - - pto::EmittedLLVMModule cubeModule; - pto::EmittedLLVMModule vectorModule; - if (failed( - pto::lowerVPTOModuleToLLVMModules(module, options, cubeModule, - vectorModule, llvm::errs()))) { - llvm::errs() << "Error: Failed to lower VPTO to LLVM modules.\n"; - return 1; - } - - if (failed(pto::emitVPTOFatobj(cubeModule.module.get(), - vectorModule.module.get(), stubSource, - outputFile, llvm::errs()))) { - llvm::errs() << "Error: Failed to emit VPTO fatobj.\n"; - return 1; - } - outputFile.keep(); - return 0; -} - -static LogicalResult runVPTOBackendPipeline(OwningOpRef &module, - int argc, char **argv, - bool hasTileOpsToExpand, - bool hasTilelangHelpers) { - PassManager pm(module->getContext()); - pm.enableVerifier(); - pm.addPass(pto::createVPTOSplitCVModulePass()); - pm.addPass(pto::createVPTONormalizeContainerPass()); - if (!hasTileOpsToExpand && hasTilelangHelpers) - inlineTilelangHelpersOnVPTOInput(pm); - if (hasTileOpsToExpand) - lowerPTOToVPTOBackend(pm, argc, argv); - prepareVPTOForEmission(pm); - if (failed(applyConfiguredPassManagerCLOptions( - pm, "VPTO unified emission pipeline"))) - return failure(); - if (failed(pm.run(module.get()))) { - llvm::errs() << "Error: VPTO emission pipeline failed.\n"; - return failure(); - } - return success(); -} +#include "PTOASDriver.h" int main(int argc, char **argv) { - DialectRegistry registry; - registry.insert(); - registry.insert(); - registry.insert(); - registry.insert(); - registry.insert(); - registry.insert(); - registry.insert(); - registry.insert(); - registry.insert(); - - registry.insert(); - arith::registerBufferizableOpInterfaceExternalModels(registry); - tensor::registerBufferizableOpInterfaceExternalModels(registry); - pto::registerBufferizableOpInterfaceExternalModels(registry); - - registry.insert(); - registry.insert(); - mlir::registerAllPasses(); - ::registerPTOPasses(); - mlir::pto::registerPTOViewToMemrefPass(); - ::registerPTOInlineLibCall(); - ::registerFoldTileBufIntrinsics(); - ::registerExpandTileOp(); - mlir::registerPassManagerCLOptions(); - - llvm::cl::SetVersionPrinter(printPTOASVersion); - - bool cliArchSpecified = false; - for (int i = 1; i < argc; ++i) { - llvm::StringRef arg(argv[i]); - if (arg == "--pto-arch" || arg.starts_with("--pto-arch=")) - cliArchSpecified = true; - } - - // Parse command line options - llvm::cl::ParseCommandLineOptions(argc, argv, "PTO Assembler (ptoas)\n"); - - PTOBackend effectiveBackend = PTOBackend::EmitC; - if (!parseBackend(ptoBackend, effectiveBackend)) { - llvm::errs() << "Error: invalid --pto-backend='" << ptoBackend - << "'. Expected 'emitc' or 'vpto'.\n"; - return 1; - } - - if (effectiveBackend != PTOBackend::VPTO && - (emitVPTO || ptoPrintSeamIR || !ptoSeamIRFile.empty())) { - llvm::errs() << "Error: VPTO-specific flags require " - "--pto-backend=vpto.\n"; - return 1; - } - - // Read whole input first (so we can auto-detect .ptobc by magic). - auto fileOrErr = llvm::MemoryBuffer::getFileOrSTDIN(inputFilename); - if (!fileOrErr) { - llvm::errs() << "Error: Could not open input file: " - << fileOrErr.getError().message() << "\n"; - return 1; - } - - MLIRContext context(registry); - // Be tolerant: ptobc decode may materialize ops from dialects that aren't - // explicitly registered/loaded in this tool yet. - context.allowUnregisteredDialects(true); - - context.getOrLoadDialect(); - context.getOrLoadDialect(); - context.getOrLoadDialect(); - context.getOrLoadDialect(); - context.getOrLoadDialect(); - context.getOrLoadDialect(); - context.getOrLoadDialect(); - context.getOrLoadDialect(); - - OwningOpRef module; - llvm::StringRef buf = (*fileOrErr)->getBuffer(); - const bool isPTOBC = (buf.size() >= 6 && std::memcmp(buf.data(), "PTOBC\0", 6) == 0); - - auto normalizeArch = [](llvm::StringRef archValue) { - std::string normalized = archValue.str(); - for (char &c : normalized) - c = static_cast(std::tolower(static_cast(c))); - return normalized; - }; - auto detectTextualModuleArch = [&](llvm::StringRef text) -> std::optional { - llvm::SmallVector matches; - llvm::Regex archRegex( - R"ptoarch("?(pto\.target_arch)"?[[:space:]]*=[[:space:]]*"([[:alpha:][:digit:]_]+)")ptoarch"); - if (!archRegex.match(text, &matches) || matches.size() < 3) - return std::nullopt; - return normalizeArch(matches[2]); - }; - - std::string arch = normalizeArch(ptoTargetArch); - if (cliArchSpecified) { - if (arch != "a3" && arch != "a5") { - llvm::errs() << "Error: invalid --pto-arch='" << ptoTargetArch - << "'. Expected 'a3' or 'a5'.\n"; - return 1; - } - } else if (!isPTOBC) { - if (auto detectedArch = detectTextualModuleArch(buf)) - arch = *detectedArch; - } - if (arch != "a3" && arch != "a5") - arch = "a3"; - - if (isPTOBC) { - // Decode PTO bytecode directly into an MLIR module. - llvm::ArrayRef bytes(reinterpret_cast(buf.data()), buf.size()); -#if defined(__cpp_exceptions) || defined(__EXCEPTIONS) - try { - module = ptobc::decodePTOBCToModule(bytes, context); - } catch (...) { - llvm::errs() << "Error: Failed to decode PTOBC.\n"; - return 1; - } -#else - module = ptobc::decodePTOBCToModule(bytes, context); -#endif - if (!module) { - llvm::errs() << "Error: Failed to decode PTOBC.\n"; - return 1; - } - } else { - // Parse textual MLIR (.pto). - llvm::SourceMgr sourceMgr; - sourceMgr.AddNewSourceBuffer(std::move(*fileOrErr), llvm::SMLoc()); - pto::ScopedPTOParserTargetArch scopedParserArch( - &context, arch == "a5" ? pto::PTOParserTargetArch::A5 - : pto::PTOParserTargetArch::A3); - module = parseSourceFile(sourceMgr, &context); - if (!module) { - llvm::errs() << "Error: Failed to parse MLIR.\n"; - return 1; - } - } - - // If the CLI explicitly requested an arch, it overrides the input module. - // Otherwise, preserve the textual module's arch when present and only fall - // back to the effective default. - if (cliArchSpecified || !module->getOperation()->hasAttr("pto.target_arch")) { - module->getOperation()->setAttr("pto.target_arch", - mlir::StringAttr::get(&context, arch)); - } - - PTOBuildLevel effectiveLevel = defaultBuildLevel(); - if (!parseBuildLevel(ptoBuildLevel, effectiveLevel)) { - llvm::errs() << "Error: invalid --pto-level='" << ptoBuildLevel - << "'. Expected 'level1', 'level2', or 'level3'.\n"; - return 1; - } - - bool invalidAutoSyncTailHint = false; - module->walk([&](mlir::func::FuncOp func) { - auto hintAttr = - func->getAttrOfType("pto.auto_sync_tail_hint"); - if (!hintAttr) - return; - - std::string normalizedHint; - if (!parseAutoSyncTailHint(hintAttr.getValue(), normalizedHint)) { - func.emitError("invalid pto.auto_sync_tail_hint '") - << hintAttr.getValue() - << "'. Expected 'barrier-all' (or 'default') or " - "'mte3-to-s-event0'."; - invalidAutoSyncTailHint = true; - return; - } - func->setAttr("pto.auto_sync_tail_hint", - mlir::StringAttr::get(&context, normalizedHint)); - }); - if (invalidAutoSyncTailHint) - return 1; - - bool hasTAssign = false; - module->walk([&](pto::TAssignOp) { hasTAssign = true; }); - - if (hasTAssign && effectiveLevel != PTOBuildLevel::Level3) { - llvm::errs() << "Error: pto.tassign is only supported when " - "--pto-level=level3.\n"; - return 1; - } - - if (hasTAssign && enableInsertSync) { - llvm::errs() << "Error: pto.tassign requires --enable-insert-sync to be " - "disabled.\n"; - return 1; - } - - int enabledAutoSyncModes = - (enableInsertSync ? 1 : 0) + (enableInjectBarrierAllSync ? 1 : 0) + - (enableGraphSyncSolver ? 1 : 0); - if (enabledAutoSyncModes > 1) { - llvm::errs() << "Error: --enable-insert-sync, " - "--enable-inject-barrier-all-sync, and " - "--enable-graph-sync-solver are mutually exclusive.\n"; - return 1; - } - if (hasTAssign && enableInjectBarrierAllSync) { - llvm::errs() << "Error: pto.tassign requires " - "--enable-inject-barrier-all-sync to be disabled.\n"; - return 1; - } - if (hasTAssign && enableGraphSyncSolver) { - llvm::errs() << "Error: pto.tassign requires --enable-graph-sync-solver " - "to be disabled.\n"; - return 1; - } - - if (effectiveLevel == PTOBuildLevel::Level3) { - bool missing = false; - module->walk([&](pto::AllocTileOp op) { - if (!op.getAddr()) { - op.emitError("requires 'addr' operand when --pto-level=level3"); - missing = true; - } - }); - if (missing) - return 1; - } else { - bool hasAddr = false; - module->walk([&](pto::AllocTileOp op) { - if (op.getAddr()) { - op.emitError( - "unexpected 'addr' operand: only supported when --pto-level=level3"); - hasAddr = true; - } - }); - if (hasAddr) - return 1; - } - - // [Fix] ToolOutputFile Usage - std::error_code ec; - llvm::ToolOutputFile outputFile(outputFilename, ec, llvm::sys::fs::OF_None); - if (ec) { - llvm::errs() << ec.message() << "\n"; - return 1; - } - - const bool hasTileOpsToExpand = hasUnexpandedTileOps(*module); - const bool hasTilelangHelpers = hasTilelangInlineHelpers(*module); - - if (effectiveBackend == PTOBackend::VPTO && !hasTileOpsToExpand) { - if (ptoPrintSeamIR || !ptoSeamIRFile.empty()) { - llvm::errs() << "Error: shared pre-backend seam IR is unavailable when " - "skipping the shared PTO-to-VPTO lowering pipeline.\n"; - return 1; - } - if (failed(runVPTOBackendPipeline(module, argc, argv, hasTileOpsToExpand, - hasTilelangHelpers))) - return 1; - return emitVPTOBackendResult(module.get(), outputFile); - } - - // Main PassManager - PassManager pm(&context); - - if (failed(applyPassManagerCLOptions(pm))) - return 1; - - pm.addNestedPass( - pto::createPTOAssignDefaultFrontendPipeIdPass()); - pm.addNestedPass( - pto::createPTOLowerFrontendPipeOpsPass()); - //pm.addNestedPass(pto::createPTOVerifyTFreePass()); - pm.addPass(pto::createPTOInferValidatePipeInitPass()); - pm.addNestedPass(pto::createLoweringSyncToPipePass()); - - if (!disableInferLayout) - pm.addNestedPass(pto::createInferPTOLayoutPass()); - pm.addNestedPass(pto::createPTOA5NormalizeTMovPass()); - pm.addPass(pto::createPTOViewToMemrefPass()); - - if (effectiveLevel != PTOBuildLevel::Level3) { - PlanMemoryOptions planMemoryOption; - planMemoryOption.memMode = MemPlanMode::LOCAL_MEM_PLAN; - planMemoryOption.enableGlobalReuse = false; - planMemoryOption.enablePrintMemoryAllocatedSize = false; - pm.addPass(pto::createPlanMemoryPass(planMemoryOption)); - } - pm.addPass(pto::createPTOResolveReservedBuffersPass()); - - // Conditionally add one automatic synchronization mode. Barrier-all is a - // conservative standalone pass; InsertSync and GraphSyncSolver are set/wait - // solvers. - if (enableInsertSync) - pm.addNestedPass(pto::createPTOInsertSyncPass()); - else if (enableInjectBarrierAllSync) - pm.addNestedPass( - pto::createPTOInjectBarrierAllSyncPass()); - else if (enableGraphSyncSolver) { - PTOGraphSyncSolverOptions graphSyncOpts; - graphSyncOpts.eventIdNumMax = graphSyncSolverEventIdMax; - pm.addNestedPass( - pto::createPTOGraphSyncSolverPass(graphSyncOpts)); - } - - llvm::raw_ostream *outputOS = &outputFile.os(); - - if (emitMlirIR) { - if (failed(pm.run(*module))) { - llvm::errs() << "Error: Pass execution failed.\n"; - return 1; - } - module->print(*outputOS); - outputFile.keep(); - return 0; - } - - // Reintroduce tile-native handles once on the shared mainline so both - // backends consume the same post-planning seam IR. - pm.addPass(pto::createPTOMaterializeTileHandlesPass()); - pm.addPass(createCSEPass()); - if (failed(applyConfiguredPassManagerCLOptions(pm, "main PTOAS pipeline"))) - return 1; - - module->getOperation()->setAttr("pto.target_arch", - mlir::StringAttr::get(&context, arch)); - - if (effectiveBackend == PTOBackend::VPTO) { - if (failed(pm.run(*module))) { - llvm::errs() << "Error: Pass execution failed.\n"; - return 1; - } - - if (ptoPrintSeamIR) { - module->print(llvm::errs()); - llvm::errs() << "\n"; - } - if (failed(emitSharedPreBackendSeamIR(*module, ptoSeamIRFile))) - return 1; - - if (failed(runVPTOBackendPipeline(module, argc, argv, hasTileOpsToExpand, - hasTilelangHelpers))) - return 1; - return emitVPTOBackendResult(module.get(), outputFile); - } - - if (arch == "a3") { - pm.addPass(pto::createEmitPTOManualPass(pto::PTOArch::A3)); - } else { - pm.addPass(pto::createEmitPTOManualPass(pto::PTOArch::A5)); - } - pm.addPass(emitc::createFormExpressionsPass()); - pm.addPass(mlir::createCSEPass()); - - if (failed(pm.run(*module))) { - llvm::errs() << "Error: Pass execution failed.\n"; - return 1; - } - - dropEmptyEmitCExpressions(module.get()); - materializeControlFlowOperands(module.get()); - if (failed(reorderEmitCFunctions(module.get()))) { - llvm::errs() << "Error: Failed to order emitted functions for C++ emission.\n"; - return 1; - } - - // Emit C++ to string, then post-process, then write to output file. - std::string cppOutput; - llvm::raw_string_ostream cppOS(cppOutput); - // CFG-style lowering (e.g. scf.while -> cf.br/cf.cond_br) may introduce - // multiple blocks, requiring variables to be declared at the top for valid - // C++ emission. - bool declareVariablesAtTop = shouldDeclareVariablesAtTop(*module); - if (failed(emitc::translateToCpp(*module, cppOS, - /*declareVariablesAtTop=*/declareVariablesAtTop))) { - llvm::errs() << "Error: Failed to emit C++.\n"; - return 1; - } - cppOS.flush(); - rewriteTileGetSetValueMarkers(cppOutput); - rewriteAsyncEventMarkers(cppOutput); - rewritePtrScalarMarkers(cppOutput); - rewriteEventIdArrayMarkers(cppOutput); - rewriteAddPtrTraceMarkers(cppOutput, emitAddPtrTrace); - rewriteScalarConstantDecls(cppOutput); - rewriteHoistedGlobalTensorDecls(cppOutput); - - *outputOS << cppOutput; - outputOS->flush(); - - outputFile.keep(); // Success, keep the file - - return 0; + return mlir::pto::runPTOASDriver(argc, argv); } From 7f61cce757cdf335b9661274cb91555256f95ce0 Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Fri, 29 May 2026 00:48:06 +0800 Subject: [PATCH 2/2] feat: backport ptoas.cpp --- include/PTO/IR/PTO.h | 3 + lib/PTO/IR/PTO.cpp | 5 + lib/PTO/Transforms/VPTOLLVMEmitter.cpp | 9 +- lib/PTO/Transforms/VPTOSplitCVModule.cpp | 14 +- .../backend_mixed_child_attrs_rejected.pto | 2 +- ...end_mixed_duplicate_abi_export_invalid.pto | 2 +- ...end_reserved_aicore_abi_suffix_invalid.pto | 21 - ...end_reserved_public_abi_suffix_invalid.pto | 21 - ...vpto_child_missing_kernel_kind_invalid.pto | 2 +- test/lit/vpto/cube_load_frac.pto | 2 +- .../tilelang_soft_vmod_backend_inline.pto | 2 +- tools/ptoas/CMakeLists.txt | 2 +- tools/ptoas/PTOASDriver.cpp | 2726 ----------------- tools/ptoas/PTOASDriver.h | 18 - tools/ptoas/VPTOHostStubEmission.cpp | 6 +- tools/ptoas/driver.cpp | 822 +++++ tools/ptoas/ptoas.cpp | 1725 ++++++++++- tools/ptoas/ptoas.h | 68 + 18 files changed, 2633 insertions(+), 2817 deletions(-) delete mode 100644 test/lit/vpto/backend_reserved_aicore_abi_suffix_invalid.pto delete mode 100644 test/lit/vpto/backend_reserved_public_abi_suffix_invalid.pto delete mode 100644 tools/ptoas/PTOASDriver.cpp delete mode 100644 tools/ptoas/PTOASDriver.h create mode 100644 tools/ptoas/driver.cpp create mode 100644 tools/ptoas/ptoas.h diff --git a/include/PTO/IR/PTO.h b/include/PTO/IR/PTO.h index 5117df15c9..27b198834a 100644 --- a/include/PTO/IR/PTO.h +++ b/include/PTO/IR/PTO.h @@ -185,6 +185,9 @@ inline constexpr llvm::StringLiteral kPTOSimtEntryAttrName = "pto.simt_entry"; /// Return true if the function carries an explicit entry marker. bool hasExplicitPTOEntryAttr(func::FuncOp func); +/// Return true if the function is marked as a PTO device kernel. +bool isPTOKernelFunction(func::FuncOp func); + /// Return true if the function should be emitted as an AICORE entry. bool isPTOEntryFunction(func::FuncOp func); diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index ed4a87f731..7e0479877a 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -2111,6 +2111,11 @@ bool mlir::pto::hasExplicitPTOEntryAttr(func::FuncOp func) { func->hasAttrOfType(kLegacyHACCEntryAttrName)); } +bool mlir::pto::isPTOKernelFunction(func::FuncOp func) { + return func && (func->hasAttrOfType("pto.kernel") || + func->hasAttrOfType("pto.aicore")); +} + static constexpr StringLiteral kEffectivePTOEntryAttrName = "pto.internal.entry"; diff --git a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp index 88876d65d8..bb42762e81 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp @@ -55,16 +55,9 @@ void attachHIVMKernelAnnotations(llvm::Module &llvmModule); namespace { -constexpr llvm::StringLiteral kPTOKernelAttrName = "pto.kernel"; -constexpr llvm::StringLiteral kLegacyPTOAICoreAttrName = "pto.aicore"; constexpr llvm::StringLiteral kVectorSuffix = "_mix_aiv"; constexpr llvm::StringLiteral kCubeSuffix = "_mix_aic"; -static bool hasVPTOKernelAttr(Operation *op) { - return op->hasAttr(kPTOKernelAttrName) || - op->hasAttr(kLegacyPTOAICoreAttrName); -} - static std::string getElementTypeFragment(Type type); static Type getElementTypeFromVectorLike(Type type); static std::optional getElementCountFromVectorLike(Type type); @@ -7784,7 +7777,7 @@ static LogicalResult renameKernelFunctionsForKernelKind(ModuleOp module, } for (func::FuncOp funcOp : module.getOps()) { - if (!hasVPTOKernelAttr(funcOp)) + if (!pto::isPTOKernelFunction(funcOp)) continue; if (funcOp.getSymName().ends_with(suffix)) continue; diff --git a/lib/PTO/Transforms/VPTOSplitCVModule.cpp b/lib/PTO/Transforms/VPTOSplitCVModule.cpp index 6cbc622b54..94a6161a91 100644 --- a/lib/PTO/Transforms/VPTOSplitCVModule.cpp +++ b/lib/PTO/Transforms/VPTOSplitCVModule.cpp @@ -26,10 +26,6 @@ using namespace mlir::pto; namespace { -static bool hasVPTOKernelAttr(Operation *op) { - return op->hasAttr("pto.kernel") || op->hasAttr("pto.aicore"); -} - static bool hasKernelKind(ModuleOp module) { return module->hasAttr(FunctionKernelKindAttr::name); } @@ -42,7 +38,7 @@ static bool hasKernelKindChildModule(ModuleOp module) { static bool hasCVSections(ModuleOp module) { bool found = false; module.walk([&](func::FuncOp funcOp) { - if (found || !hasVPTOKernelAttr(funcOp)) + if (found || !pto::isPTOKernelFunction(funcOp)) return WalkResult::advance(); WalkResult result = funcOp.walk([&](Operation *op) { if (isa(op)) { @@ -60,7 +56,7 @@ static bool hasCVSections(ModuleOp module) { static bool hasSectionKind(ModuleOp module, FunctionKernelKind kind) { bool found = false; module.walk([&](func::FuncOp funcOp) { - if (found || !hasVPTOKernelAttr(funcOp)) + if (found || !pto::isPTOKernelFunction(funcOp)) return WalkResult::advance(); WalkResult result = funcOp.walk([&](Operation *op) { bool matches = kind == FunctionKernelKind::Cube @@ -125,7 +121,7 @@ static LogicalResult verifyNoNestedSections(ModuleOp module) { static LogicalResult verifyKernelFunctionsUseSections(ModuleOp module) { LogicalResult status = success(); module.walk([&](func::FuncOp funcOp) { - if (failed(status) || !hasVPTOKernelAttr(funcOp)) + if (failed(status) || !pto::isPTOKernelFunction(funcOp)) return WalkResult::advance(); if (!hasAnySection(funcOp)) { status = funcOp.emitOpError( @@ -141,7 +137,7 @@ static LogicalResult verifyKernelFunctionsUseSections(ModuleOp module) { static LogicalResult verifyUniqueSectionKindsPerFunction(ModuleOp module) { LogicalResult status = success(); module.walk([&](func::FuncOp funcOp) { - if (failed(status) || !hasVPTOKernelAttr(funcOp)) + if (failed(status) || !pto::isPTOKernelFunction(funcOp)) return WalkResult::advance(); unsigned cubeCount = 0; unsigned vectorCount = 0; @@ -168,7 +164,7 @@ static void eraseKernelFunctionsWithoutSectionKind(ModuleOp module, FunctionKernelKind kind) { SmallVector eraseFuncs; module.walk([&](func::FuncOp funcOp) { - if (hasVPTOKernelAttr(funcOp) && !hasSectionKind(funcOp, kind)) + if (pto::isPTOKernelFunction(funcOp) && !hasSectionKind(funcOp, kind)) eraseFuncs.push_back(funcOp); }); diff --git a/test/lit/vpto/backend_mixed_child_attrs_rejected.pto b/test/lit/vpto/backend_mixed_child_attrs_rejected.pto index 8367461871..a2c929e75a 100644 --- a/test/lit/vpto/backend_mixed_child_attrs_rejected.pto +++ b/test/lit/vpto/backend_mixed_child_attrs_rejected.pto @@ -25,4 +25,4 @@ module attributes {pto.target_arch = "a5"} { } } -// CHECK: mixed pto.backend fatobj mode does not support debug IR output flags +// CHECK: mixed pto.backend fatobj mode requires an explicit file path passed with -o diff --git a/test/lit/vpto/backend_mixed_duplicate_abi_export_invalid.pto b/test/lit/vpto/backend_mixed_duplicate_abi_export_invalid.pto index d1bfdf21c5..9ea92f43c5 100644 --- a/test/lit/vpto/backend_mixed_duplicate_abi_export_invalid.pto +++ b/test/lit/vpto/backend_mixed_duplicate_abi_export_invalid.pto @@ -6,7 +6,7 @@ // 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=a5 --emit-vpto %s -o - 2>&1 | FileCheck %s +// RUN: not ptoas --pto-arch=a5 %s -o %t.o 2>&1 | FileCheck %s module attributes {pto.target_arch = "a5"} { module attributes {pto.backend = "emitc"} { diff --git a/test/lit/vpto/backend_reserved_aicore_abi_suffix_invalid.pto b/test/lit/vpto/backend_reserved_aicore_abi_suffix_invalid.pto deleted file mode 100644 index c57562a709..0000000000 --- a/test/lit/vpto/backend_reserved_aicore_abi_suffix_invalid.pto +++ /dev/null @@ -1,21 +0,0 @@ -// 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: not ptoas --pto-arch=a5 --emit-vpto %s -o - 2>&1 | FileCheck %s - -module attributes { - pto.backend = "vpto", - pto.kernel_kind = #pto.kernel_kind, - pto.target_arch = "a5" -} { - func.func @device_kernel_mix_aic() attributes {pto.aicore} { - return - } -} - -// CHECK: source-level pto.aicore function 'device_kernel_mix_aic' must not use reserved device ABI suffix '_mix_aic' diff --git a/test/lit/vpto/backend_reserved_public_abi_suffix_invalid.pto b/test/lit/vpto/backend_reserved_public_abi_suffix_invalid.pto deleted file mode 100644 index 13a03fe5b0..0000000000 --- a/test/lit/vpto/backend_reserved_public_abi_suffix_invalid.pto +++ /dev/null @@ -1,21 +0,0 @@ -// 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: not ptoas --pto-arch=a5 --emit-vpto %s -o - 2>&1 | FileCheck %s - -module attributes { - pto.backend = "vpto", - pto.kernel_kind = #pto.kernel_kind, - pto.target_arch = "a5" -} { - func.func public @vpto_post.vector() { - return - } -} - -// CHECK: source-level public non-pto.aicore function 'vpto_post.vector' must not use reserved VPTO ABI suffix '.vector' diff --git a/test/lit/vpto/backend_vpto_child_missing_kernel_kind_invalid.pto b/test/lit/vpto/backend_vpto_child_missing_kernel_kind_invalid.pto index 425925de04..9f8892de01 100644 --- a/test/lit/vpto/backend_vpto_child_missing_kernel_kind_invalid.pto +++ b/test/lit/vpto/backend_vpto_child_missing_kernel_kind_invalid.pto @@ -16,4 +16,4 @@ module attributes {pto.target_arch = "a5"} { } } -// CHECK: vpto backend child module must carry 'pto.kernel_kind' +// CHECK: expected VPTO kernel submodule to carry 'pto.kernel_kind' diff --git a/test/lit/vpto/cube_load_frac.pto b/test/lit/vpto/cube_load_frac.pto index 7699f742e9..a9335eae15 100644 --- a/test/lit/vpto/cube_load_frac.pto +++ b/test/lit/vpto/cube_load_frac.pto @@ -48,7 +48,7 @@ module attributes {"pto.target_arch" = "a5", pto.kernel_kind = #pto.kernel_kind< // ROUNDTRIP: pto.mte_gm_l1_frac %{{.*}}, %{{.*}}, nd2nz, shape(%{{.*}}, %{{.*}}), src_layout(%{{.*}}, %{{.*}}), dst_group(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}), ctrl(%{{.*}}, %{{.*}}) // ROUNDTRIP: pto.mte_gm_l1_frac %{{.*}}, %{{.*}}, dn2nz, shape(%{{.*}}, %{{.*}}), src_layout(%{{.*}}), dst_group(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}), ctrl(%{{.*}}, %{{.*}}) -// EXPAND-LABEL: func.func @cube_load_frac_roundtrip( +// EXPAND-LABEL: func.func @cube_load_frac_roundtrip.cube( // EXPAND: pto.set_mte2_nz_para // EXPAND: pto.copy_gm_to_cbuf_multi_nd2nz // EXPAND: pto.set_mte2_nz_para diff --git a/test/lit/vpto/tilelang_soft_vmod_backend_inline.pto b/test/lit/vpto/tilelang_soft_vmod_backend_inline.pto index 65c84c2ec1..fb8613ab9d 100644 --- a/test/lit/vpto/tilelang_soft_vmod_backend_inline.pto +++ b/test/lit/vpto/tilelang_soft_vmod_backend_inline.pto @@ -112,7 +112,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind -#include -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/Tensor/IR/Tensor.h" -#include "mlir/Target/Cpp/CppEmitter.h" -#include "llvm/Support/SourceMgr.h" -#include "llvm/Support/ToolOutputFile.h" -#include "llvm/Support/FileSystem.h" // [Fix] Required for OF_None -#include "llvm/Support/Path.h" -#include "ptobc/ptobc_decode.h" -#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h" -#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" -#include "mlir/Dialect/LLVMIR/LLVMDialect.h" -#include "mlir/Dialect/EmitC/IR/EmitC.h" -#include "mlir/Dialect/EmitC/Transforms/Passes.h" -#include "llvm/Support/CommandLine.h" -#include "llvm/Support/Regex.h" -#include "llvm/Support/raw_ostream.h" -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/ScopeExit.h" -#include "llvm/ADT/SmallPtrSet.h" -#include "llvm/ADT/SmallVector.h" -#include "llvm/ADT/StringRef.h" -#include "llvm/ADT/StringSwitch.h" -#include "llvm/ADT/StringMap.h" -#include "llvm/Support/MemoryBuffer.h" -#include "llvm/Support/Process.h" -#include "llvm/Support/Program.h" -#include -#include -#include -#include -#include -#include -#include - -extern "C" { -extern char **environ; -} - -using namespace mlir; -using namespace pto; - -#ifndef PTOAS_RELEASE_VERSION -#define PTOAS_RELEASE_VERSION "unknown" -#endif - -extern "C" int main(int argc, char **argv); - -static void printPTOASVersion(llvm::raw_ostream &os) { - os << "ptoas " << PTOAS_RELEASE_VERSION << "\n"; -} - -static std::string getParentDir(llvm::StringRef path) { - llvm::SmallString<256> parent(path); - llvm::sys::path::remove_filename(parent); - llvm::sys::path::remove_dots(parent, true); - return std::string(parent); -} - -static bool pathExists(llvm::StringRef path) { - return !path.empty() && llvm::sys::fs::exists(path); -} - -static std::string joinPath(llvm::StringRef lhs, llvm::StringRef rhs) { - llvm::SmallString<256> joined(lhs); - llvm::sys::path::append(joined, rhs); - llvm::sys::path::remove_dots(joined, true); - return std::string(joined); -} - -static std::string detectInstalledTilelangPath(const char *argv0) { - std::string exePath = llvm::sys::fs::getMainExecutable(argv0, (void *)&main); - if (exePath.empty()) - return {}; - - const std::string exeDir = getParentDir(exePath); - const std::string prefixDir = getParentDir(exeDir); - const std::string installedTileOps = joinPath(prefixDir, "share/ptoas/TileOps"); - if (pathExists(installedTileOps)) - return installedTileOps; - return {}; -} - -static std::string detectInstalledTilelangPkgPath(const char *argv0) { - std::string exePath = llvm::sys::fs::getMainExecutable(argv0, (void *)&main); - if (exePath.empty()) - return {}; - - const std::string exeDir = getParentDir(exePath); - const std::string prefixDir = getParentDir(exeDir); - const std::string installedPkgRoot = prefixDir; - const std::string installedPkg = joinPath(installedPkgRoot, "tilelang_dsl"); - if (pathExists(installedPkg)) - return installedPkgRoot; - return {}; -} - -static bool hasCLIOption(int argc, char **argv, llvm::StringRef option) { - const std::string optionWithValue = (option + "=").str(); - for (int i = 1; i < argc; ++i) { - llvm::StringRef arg(argv[i]); - if (arg == option || arg.starts_with(optionWithValue)) - return true; - } - return false; -} - -static LogicalResult applyConfiguredPassManagerCLOptions( - PassManager &pm, llvm::StringRef pipelineName, - llvm::raw_ostream &diagOS = llvm::errs()) { - if (succeeded(mlir::applyPassManagerCLOptions(pm))) - return success(); - diagOS << "Error: failed to apply MLIR pass manager command-line options for " - << pipelineName << ".\n"; - return failure(); -} - -static LogicalResult reorderEmitCFunctions(ModuleOp module) { - SmallVector declarations; - SmallVector definitions; - llvm::DenseMap definitionsByName; - - for (auto func : module.getOps()) { - if (func.isDeclaration()) { - declarations.push_back(func); - continue; - } - definitions.push_back(func); - definitionsByName[func.getSymNameAttr()] = func; - } - - llvm::DenseMap indegree; - llvm::DenseMap> outgoing; - for (auto func : definitions) - indegree[func.getOperation()] = 0; - - for (auto caller : definitions) { - Operation *callerOp = caller.getOperation(); - llvm::SmallPtrSet seenCallees; - bool hasCycle = false; - caller.walk([&](emitc::CallOp call) { - auto calleeAttr = call.getCalleeAttr(); - if (!calleeAttr) - return; - auto it = definitionsByName.find(calleeAttr.getLeafReference()); - if (it == definitionsByName.end()) - return; - Operation *calleeOp = it->second.getOperation(); - if (calleeOp == callerOp) { - hasCycle = true; - return; - } - if (!seenCallees.insert(calleeOp).second) - return; - outgoing[calleeOp].push_back(callerOp); - ++indegree[callerOp]; - }); - if (hasCycle) { - return caller.emitOpError() - << "recursive function calls are not supported for EmitC C++ " - "emission"; - } - } - - SmallVector ready; - for (auto func : definitions) { - if (indegree[func.getOperation()] == 0) - ready.push_back(func.getOperation()); - } - - SmallVector sortedDefinitions; - while (!ready.empty()) { - Operation *next = ready.front(); - ready.erase(ready.begin()); - auto nextFunc = cast(next); - sortedDefinitions.push_back(nextFunc); - - for (Operation *user : outgoing[next]) { - unsigned &userIndegree = indegree[user]; - if (--userIndegree == 0) - ready.push_back(user); - } - } - - if (sortedDefinitions.size() != definitions.size()) { - return module.emitError() - << "cyclic function call graph is not supported for EmitC C++ emission"; - } - - if (declarations.empty() && definitions.size() <= 1) - return success(); - - SmallVector desiredOrder; - desiredOrder.append(declarations.begin(), declarations.end()); - desiredOrder.append(sortedDefinitions.begin(), sortedDefinitions.end()); - - Block &body = module.getBodyRegion().front(); - Operation *anchor = nullptr; - for (Operation &op : body.getOperations()) { - if (isa(op)) { - anchor = &op; - break; - } - } - if (!anchor) - return success(); - - auto advanceAnchor = [&]() { - while (anchor) { - anchor = anchor->getNextNode(); - if (!anchor || isa(anchor)) - return; - } - }; - - for (auto func : desiredOrder) { - if (func.getOperation() == anchor) { - advanceAnchor(); - continue; - } - if (anchor) - func->moveBefore(anchor); - else - func->moveBefore(&body, body.end()); - } - - return success(); -} - -// -------------------------------------------------------------------------- -// Command Line Options -// -------------------------------------------------------------------------- -static llvm::cl::opt inputFilename(llvm::cl::Positional, - llvm::cl::desc(""), - llvm::cl::init("-")); - -static llvm::cl::opt outputFilename("o", - llvm::cl::desc("Output filename"), - llvm::cl::value_desc("filename"), - llvm::cl::init("-")); - -static llvm::cl::opt enableInsertSync("enable-insert-sync", - llvm::cl::desc("Enable automatic synchronization insertion pass"), - llvm::cl::init(false)); - -static llvm::cl::opt enableInjectBarrierAllSync( - "enable-inject-barrier-all-sync", - llvm::cl::desc("Enable conservative synchronization by inserting " - "pto.barrier PIPE_ALL before memory-effecting PTO pipe ops"), - llvm::cl::init(false)); - -static llvm::cl::opt enableGraphSyncSolver( - "enable-graph-sync-solver", - llvm::cl::desc("Enable the graph-based intra-core sync solver " - "(experimental). Mutually exclusive with " - "--enable-insert-sync and " - "--enable-inject-barrier-all-sync."), - llvm::cl::init(false)); - -static llvm::cl::opt graphSyncSolverEventIdMax( - "graph-sync-solver-event-id-max", - llvm::cl::desc( - "Maximum EVENT_ID slots for the graph sync solver (default 8). " - "Lower values exercise the PIPE_ALL coloring fallback sooner."), - llvm::cl::init(8)); - -static llvm::cl::opt enableTileOpExpand( - "enable-tile-op-expand", - llvm::cl::desc( - "Deprecated compatibility flag. TileOp expansion is controlled by " - "--pto-backend=vpto."), - llvm::cl::init(false)); - -#ifndef PTOAS_DEFAULT_TILELANG_PATH -#define PTOAS_DEFAULT_TILELANG_PATH "" -#endif -#ifndef PTOAS_DEFAULT_TILELANG_PKG_PATH -#define PTOAS_DEFAULT_TILELANG_PKG_PATH "" -#endif - -static llvm::cl::opt tilelangPath( - "tilelang-path", - llvm::cl::desc("Path to directory of .py tilelang DSL template files " - "(default: /lib/TileOps, baked in at build time)"), - llvm::cl::init(PTOAS_DEFAULT_TILELANG_PATH)); - -static llvm::cl::opt tilelangPkgPath( - "tilelang-pkg-path", - llvm::cl::desc("PYTHONPATH for tilelang_dsl package " - "(default: /tilelang-dsl/python, baked in at build time)"), - llvm::cl::init(PTOAS_DEFAULT_TILELANG_PKG_PATH)); - -static llvm::cl::opt daemonSocketPath( - "daemon-socket-path", - llvm::cl::desc("Path to Unix domain socket for daemon RPC " - "(default: /tmp/tilelang_daemon_{pid}.sock)"), - llvm::cl::init("")); - -static pto::ExpandTileOpOptions resolveExpandTileOpOptions(int argc, - char **argv) { - pto::ExpandTileOpOptions expandOpts; - expandOpts.tilelangPath = tilelangPath; - expandOpts.tilelangPkgPath = tilelangPkgPath; - - if (!hasCLIOption(argc, argv, "--tilelang-path")) { - std::string detectedTilelangPath = detectInstalledTilelangPath(argv[0]); - if (!detectedTilelangPath.empty()) - expandOpts.tilelangPath = detectedTilelangPath; - } - - if (!hasCLIOption(argc, argv, "--tilelang-pkg-path")) { - std::string detectedTilelangPkgPath = detectInstalledTilelangPkgPath(argv[0]); - if (!detectedTilelangPkgPath.empty()) - expandOpts.tilelangPkgPath = detectedTilelangPkgPath; - } - - // Daemon mode is default (no CLI option needed) - // Automatically start daemon for instance caching - if (!expandOpts.tilelangPath.empty()) { - std::string socket = daemonSocketPath; - if (socket.empty()) - socket = ptoas::DaemonManager::generateSocketPath(); - - // Register cleanup handler (daemon will be stopped on PTOAS exit) - ptoas::registerDaemonCleanup(); - - // Try to start daemon automatically - if (ptoas::DaemonManager::start(socket, expandOpts.tilelangPath, expandOpts.tilelangPkgPath)) { - expandOpts.daemonSocketPath = socket; - llvm::errs() << "Info: TileLang daemon started successfully\n"; - } else { - // Fallback: daemon failed, use subprocess mode (current approach) - expandOpts.daemonSocketPath = ""; - llvm::errs() << "Warning: Failed to start daemon, using subprocess mode (fallback)\n"; - } - } - - return expandOpts; -} - -static llvm::cl::opt disableInferLayout( - "disable-infer-layout", - llvm::cl::desc("Disable PTO layout inference pass (static-only)"), - llvm::cl::init(false)); - -static llvm::cl::opt emitAddPtrTrace( - "emit-addptr-trace", - llvm::cl::desc("Emit addptr trace comments in generated C++ output"), - llvm::cl::init(false)); - -static llvm::cl::opt emitMlirIR( - "emit-pto-ir", - llvm::cl::desc("Emit PTO IR after lowering instead of C++"), - llvm::cl::init(false)); - -static llvm::cl::opt ptoTargetArch( - "pto-arch", - llvm::cl::desc("Target Ascend architecture for codegen: a3 or a5 (default: a3)"), - llvm::cl::value_desc("a3|a5"), - llvm::cl::init("a3")); - -static llvm::cl::opt ptoBuildLevel( - "pto-level", - llvm::cl::desc("Build level for pass pipeline: level1, level2, or level3 (default: level2)"), - llvm::cl::value_desc("level1|level2|level3"), - llvm::cl::init("level2")); - -static llvm::cl::opt ptoBackend( - "pto-backend", - llvm::cl::desc("Final PTOAS backend: emitc or vpto (default: emitc)"), - llvm::cl::value_desc("emitc|vpto"), llvm::cl::init("emitc")); - -static llvm::cl::opt emitVPTO( - "emit-vpto", - llvm::cl::desc("Write final post-pass VPTO IR to -o"), - llvm::cl::init(false)); - -static llvm::cl::opt vptoPrintIR( - "vpto-print-ir", - llvm::cl::desc("Print post-pass VPTO backend IR to stderr"), - llvm::cl::init(false)); - -static llvm::cl::opt vptoLoweringStrategy( - "vpto-lowering-strategy", - llvm::cl::desc("VPTO vector lowering strategy: post-update or no-post-update"), - llvm::cl::value_desc("post-update|no-post-update"), - llvm::cl::init("post-update")); - -static llvm::cl::opt dumpVPTOIR( - "dump-vpto-ir", - llvm::cl::desc("Print post-pass VPTO backend IR to stderr"), - llvm::cl::init(false)); - -static llvm::cl::opt ptoPrintSeamIR( - "pto-print-seam-ir", - llvm::cl::desc("Print shared pre-backend seam IR to stderr"), - llvm::cl::init(false)); - -static llvm::cl::opt ptoSeamIRFile( - "pto-seam-ir-file", - llvm::cl::desc("Write shared pre-backend seam IR to a file"), - llvm::cl::value_desc("path"), - llvm::cl::init("")); - -enum class PTOBuildLevel { - Level1, - Level2, - Level3, -}; - -enum class PTOBackend { - EmitC, - VPTO, -}; - -static PTOBuildLevel defaultBuildLevel() { - return PTOBuildLevel::Level2; -} - -static bool parseBuildLevel(llvm::StringRef levelStr, PTOBuildLevel &out) { - std::string s = levelStr.str(); - for (char &c : s) - c = static_cast(std::tolower(static_cast(c))); - if (s == "level1") { - out = PTOBuildLevel::Level1; - return true; - } - if (s == "level2") { - out = PTOBuildLevel::Level2; - return true; - } - if (s == "level3") { - out = PTOBuildLevel::Level3; - return true; - } - return false; -} - -static constexpr llvm::StringLiteral kAutoSyncTailPolicyBarrierAll = - "barrier_all"; -static constexpr llvm::StringLiteral kAutoSyncTailPolicyMte3ToSEvent0 = - "setwait_mte3_to_s_event0"; - -static bool parseAutoSyncTailHint(llvm::StringRef hintStr, std::string &normalized) { - std::string s = hintStr.str(); - for (char &c : s) - c = static_cast(std::tolower(static_cast(c))); - if (s == "barrier-all" || s == "barrier_all" || s == "default") { - normalized = kAutoSyncTailPolicyBarrierAll.str(); - return true; - } - if (s == "mte3-to-s-event0" || s == "mte3_to_s_event0" || - s == "setwait-mte3-to-s-event0" || - s == "setwait_mte3_to_s_event0") { - normalized = kAutoSyncTailPolicyMte3ToSEvent0.str(); - return true; - } - return false; -} - -static bool parseBackend(llvm::StringRef backendStr, PTOBackend &out) { - std::string s = backendStr.str(); - for (char &c : s) - c = static_cast(std::tolower(static_cast(c))); - if (s == "emitc") { - out = PTOBackend::EmitC; - return true; - } - if (s == "vpto") { - out = PTOBackend::VPTO; - return true; - } - return false; -} - -static LogicalResult parsePTOBackendAttr(Operation *op, - std::optional &backend) { - backend = std::nullopt; - Attribute rawBackendAttr = op->getAttr("pto.backend"); - if (!rawBackendAttr) - return success(); - - auto backendAttr = dyn_cast(rawBackendAttr); - if (!backendAttr) { - return op->emitError("invalid pto.backend attribute. Expected string " - "value 'emitc' or 'vpto'."); - } - - PTOBackend attrBackend = PTOBackend::EmitC; - if (!parseBackend(backendAttr.getValue(), attrBackend)) { - return op->emitError("invalid pto.backend '") - << backendAttr.getValue() << "'. Expected 'emitc' or 'vpto'."; - } - - backend = attrBackend; - return success(); -} - -static LogicalResult validatePTOBackendAttrs(ModuleOp module) { - LogicalResult status = success(); - module.walk([&](ModuleOp nested) { - if (failed(status)) - return WalkResult::interrupt(); - std::optional backend; - if (failed(parsePTOBackendAttr(nested.getOperation(), backend))) { - status = failure(); - return WalkResult::interrupt(); - } - return WalkResult::advance(); - }); - return status; -} - -static LogicalResult verifyBackendChildShape(ModuleOp module, - bool cliBackendSpecified) { - if (cliBackendSpecified || module->hasAttr("pto.backend")) - return success(); - - bool sawEmitC = false; - bool sawVPTO = false; - SmallVector missingBackendChildren; - for (ModuleOp child : module.getOps()) { - std::optional childBackend; - if (failed(parsePTOBackendAttr(child.getOperation(), childBackend))) - return failure(); - if (!childBackend) { - missingBackendChildren.push_back(child); - continue; - } - sawEmitC |= *childBackend == PTOBackend::EmitC; - sawVPTO |= *childBackend == PTOBackend::VPTO; - } - - if (!sawEmitC || !sawVPTO) - return success(); - - if (!missingBackendChildren.empty()) { - return missingBackendChildren.front().emitError() - << "mixed-backend container child module is missing pto.backend"; - } - - return success(); -} - -static LogicalResult resolveModuleBackend(ModuleOp module, - bool cliBackendSpecified, - bool &mixedBackendMode, - PTOBackend &effectiveBackend) { - mixedBackendMode = false; - if (cliBackendSpecified) - return success(); - - std::optional topBackend; - if (failed(parsePTOBackendAttr(module.getOperation(), topBackend))) - return failure(); - if (topBackend) { - effectiveBackend = *topBackend; - return success(); - } - - std::optional childBackend; - bool sawEmitCChild = false; - bool sawVPTOChild = false; - SmallVector missingBackendChildren; - for (ModuleOp child : module.getOps()) { - std::optional parsedChildBackend; - if (failed(parsePTOBackendAttr(child.getOperation(), parsedChildBackend))) - return failure(); - if (!parsedChildBackend) { - missingBackendChildren.push_back(child); - continue; - } - sawEmitCChild |= *parsedChildBackend == PTOBackend::EmitC; - sawVPTOChild |= *parsedChildBackend == PTOBackend::VPTO; - - if (!childBackend) { - childBackend = *parsedChildBackend; - continue; - } - if (*childBackend != *parsedChildBackend) { - if (!missingBackendChildren.empty()) { - return missingBackendChildren.front().emitError() - << "mixed-backend container child module is missing pto.backend"; - } - mixedBackendMode = true; - continue; - } - } - - if (childBackend) - effectiveBackend = *childBackend; - if (sawEmitCChild && sawVPTOChild && !missingBackendChildren.empty()) { - return missingBackendChildren.front().emitError() - << "mixed-backend container child module is missing pto.backend"; - } - return success(); -} - -static bool hasFunctionKernelKind(ModuleOp module) { - return static_cast(module->getAttrOfType( - FunctionKernelKindAttr::name)); -} - -static llvm::StringRef getPublicABISuffix(llvm::StringRef symName) { - if (symName.ends_with(".vector")) - return ".vector"; - if (symName.ends_with(".cube")) - return ".cube"; - return {}; -} - -static llvm::StringRef getDeviceABISuffix(llvm::StringRef symName) { - if (symName.ends_with("_mix_aiv")) - return "_mix_aiv"; - if (symName.ends_with("_mix_aic")) - return "_mix_aic"; - return {}; -} - -static LogicalResult verifyVPTOSourceSymbolNames(ModuleOp module) { - LogicalResult status = success(); - module.walk([&](func::FuncOp func) { - if (failed(status)) - return WalkResult::interrupt(); - - llvm::StringRef symName = func.getSymName(); - if (func->hasAttr("pto.aicore")) { - llvm::StringRef suffix = getDeviceABISuffix(symName); - if (!suffix.empty()) { - status = func.emitError() - << "source-level pto.aicore function '" << symName - << "' must not use reserved device ABI suffix '" << suffix - << "'; PTOAS adds it from pto.kernel_kind"; - return WalkResult::interrupt(); - } - return WalkResult::advance(); - } - - if (!func.isPublic()) - return WalkResult::advance(); - - llvm::StringRef suffix = getPublicABISuffix(symName); - if (!suffix.empty()) { - status = func.emitError() - << "source-level public non-pto.aicore function '" << symName - << "' must not use reserved VPTO ABI suffix '" << suffix - << "'; PTOAS adds it from pto.kernel_kind"; - return WalkResult::interrupt(); - } - return WalkResult::advance(); - }); - return status; -} - -static LogicalResult verifyVPTOBackendChild(ModuleOp child) { - if (!hasFunctionKernelKind(child)) { - return child.emitError() - << "vpto backend child module must carry 'pto.kernel_kind'"; - } - return verifyVPTOSourceSymbolNames(child); -} - -struct BackendExportPlan { - std::string sourceName; - std::string abiName; - FunctionType type; -}; - -struct BackendImportPlan { - std::string sourceName; - std::string abiName; - FunctionType type; -}; - -struct BackendChildPlan { - ModuleOp sourceModule; - ModuleOp module; - OwningOpRef ownedModule; - PTOBackend backend = PTOBackend::EmitC; - std::optional kernelKind; - SmallVector exports; - SmallVector imports; - SmallVector sourceDeviceSymbols; - SmallVector deviceABISymbols; -}; - -enum class CompilationMode { - EmitC, - VPTO, - MixedFatobj, -}; - -enum class BackendArtifactKind { - VectorObject, - CubeObject, - GenericObject, - Fatobj, -}; - -struct BackendArtifact { - PTOBackend backend = PTOBackend::EmitC; - BackendArtifactKind kind = BackendArtifactKind::GenericObject; - std::string path; - SmallVector exportedABISymbols; -}; - -struct BackendFatobjGroup { - PTOBackend backend = PTOBackend::EmitC; - BackendArtifactKind kind = BackendArtifactKind::GenericObject; - SmallVector objectPaths; -}; - -struct CompilationPlan { - CompilationMode mode = CompilationMode::EmitC; - PTOBackend effectiveBackend = PTOBackend::EmitC; - SmallVector children; -}; - -static std::optional getModuleKernelKind(ModuleOp module) { - auto kernelKind = module->getAttrOfType( - FunctionKernelKindAttr::name); - if (!kernelKind) - return std::nullopt; - return kernelKind.getKernelKind(); -} - -static llvm::StringRef getPublicABISuffix(FunctionKernelKind kind) { - switch (kind) { - case FunctionKernelKind::Vector: - return ".vector"; - case FunctionKernelKind::Cube: - return ".cube"; - } - llvm_unreachable("unknown function kernel kind"); -} - -static llvm::StringRef getDeviceABISuffix(FunctionKernelKind kind) { - switch (kind) { - case FunctionKernelKind::Vector: - return "_mix_aiv"; - case FunctionKernelKind::Cube: - return "_mix_aic"; - } - llvm_unreachable("unknown function kernel kind"); -} - -static void collectBackendChildSymbols(BackendChildPlan &plan) { - for (func::FuncOp func : plan.module.getOps()) { - llvm::StringRef symName = func.getSymName(); - if (func->hasAttr("pto.aicore")) { - plan.sourceDeviceSymbols.push_back(symName.str()); - if (plan.backend == PTOBackend::VPTO && plan.kernelKind) { - plan.deviceABISymbols.push_back( - (symName + getDeviceABISuffix(*plan.kernelKind)).str()); - } - continue; - } - - if (func.isDeclaration()) { - BackendImportPlan import; - import.sourceName = symName.str(); - import.type = func.getFunctionType(); - plan.imports.push_back(std::move(import)); - continue; - } - - if (!func.isPublic()) - continue; - - BackendExportPlan exportPlan; - exportPlan.sourceName = symName.str(); - exportPlan.type = func.getFunctionType(); - if (plan.backend == PTOBackend::VPTO && plan.kernelKind) { - exportPlan.abiName = - (symName + getPublicABISuffix(*plan.kernelKind)).str(); - } else { - exportPlan.abiName = symName.str(); - } - plan.exports.push_back(std::move(exportPlan)); - } -} - -static OwningOpRef cloneBackendChildModule(ModuleOp outer, - ModuleOp child) { - Operation *clonedOp = child.getOperation()->clone(); - auto clonedModule = cast(clonedOp); - - for (NamedAttribute attr : outer->getAttrs()) { - StringRef attrName = attr.getName().getValue(); - if (attrName == SymbolTable::getSymbolAttrName() || attrName == "pto.backend") - continue; - if (!clonedModule->hasAttr(attr.getName())) - clonedModule->setAttr(attr.getName(), attr.getValue()); - } - - return OwningOpRef(clonedModule); -} - -static LogicalResult -collectBackendChildren(ModuleOp module, - SmallVectorImpl &children) { - for (ModuleOp child : module.getOps()) { - std::optional childBackend; - if (failed(parsePTOBackendAttr(child.getOperation(), childBackend))) - return failure(); - if (!childBackend) - continue; - BackendChildPlan plan; - plan.sourceModule = child; - plan.ownedModule = cloneBackendChildModule(module, child); - plan.module = plan.ownedModule.get(); - plan.backend = *childBackend; - plan.kernelKind = getModuleKernelKind(plan.module); - collectBackendChildSymbols(plan); - children.push_back(std::move(plan)); - } - return success(); -} - -static CompilationMode getCompilationMode(PTOBackend effectiveBackend, - ArrayRef children) { - bool sawEmitC = false; - bool sawVPTO = false; - for (const BackendChildPlan &child : children) { - sawEmitC |= child.backend == PTOBackend::EmitC; - sawVPTO |= child.backend == PTOBackend::VPTO; - } - if (sawEmitC && sawVPTO) - return CompilationMode::MixedFatobj; - return effectiveBackend == PTOBackend::VPTO ? CompilationMode::VPTO - : CompilationMode::EmitC; -} - -static LogicalResult verifyUniqueBackendABIExports( - SmallVectorImpl &children) { - llvm::StringMap seen; - for (BackendChildPlan &child : children) { - for (const BackendExportPlan &exportPlan : child.exports) { - auto inserted = seen.try_emplace(exportPlan.abiName, child.module.getLoc()); - if (!inserted.second) { - return child.module.emitError() - << "generated ABI export symbol '" << exportPlan.abiName - << "' is defined by multiple backend child modules"; - } - } - } - return success(); -} - -static LogicalResult verifyExternalImportShape( - SmallVectorImpl &children) { - for (BackendChildPlan &child : children) { - for (func::FuncOp func : child.module.getOps()) { - if (!func.isDeclaration()) - continue; - if (!func.isPrivate()) { - return func.emitError() - << "external func.func declaration must use private visibility"; - } - if (func->hasAttr("pto.aicore")) { - return func.emitError() - << "external func.func declaration must not be marked pto.aicore"; - } - } - } - return success(); -} - -static LogicalResult resolveExternalImports( - SmallVectorImpl &children) { - for (BackendChildPlan &importingChild : children) { - for (BackendImportPlan &importPlan : importingChild.imports) { - const BackendExportPlan *matchedExport = nullptr; - const BackendChildPlan *matchedChild = nullptr; - for (const BackendChildPlan &exportingChild : children) { - if (&exportingChild == &importingChild) - continue; - for (const BackendExportPlan &exportPlan : exportingChild.exports) { - if (exportPlan.sourceName != importPlan.sourceName) - continue; - if (matchedExport) { - return importingChild.module.emitError() - << "cross-backend external import '" << importPlan.sourceName - << "' matches multiple exported definitions"; - } - matchedExport = &exportPlan; - matchedChild = &exportingChild; - } - } - - if (!matchedExport) { - return importingChild.module.emitError() - << "cross-backend external import '" << importPlan.sourceName - << "' has no matching public non-pto.aicore definition in " - "another backend child module"; - } - - if (matchedExport->type != importPlan.type) { - return importingChild.module.emitError() - << "cross-backend external import '" << importPlan.sourceName - << "' signature does not match exported definition"; - } - - importPlan.abiName = matchedExport->abiName; - (void)matchedChild; - } - } - return success(); -} - -static LogicalResult applyResolvedExternalImports(BackendChildPlan &child) { - llvm::StringMap importABIMap; - for (const BackendImportPlan &importPlan : child.imports) { - if (importPlan.abiName.empty()) - return child.module.emitError() - << "unresolved cross-backend external import '" - << importPlan.sourceName << "'"; - importABIMap[importPlan.sourceName] = importPlan.abiName; - } - if (importABIMap.empty()) - return success(); - - for (func::FuncOp func : child.module.getOps()) { - auto it = importABIMap.find(func.getSymName()); - if (it == importABIMap.end()) - continue; - func->setAttr("pto.external_abi", - StringAttr::get(func.getContext(), it->second)); - } - return success(); -} - -static LogicalResult applyVPTOPublicABIExportNames(ModuleOp module) { - LogicalResult status = success(); - module.walk([&](ModuleOp child) { - if (failed(status)) - return WalkResult::interrupt(); - std::optional kernelKind = getModuleKernelKind(child); - if (!kernelKind) - return WalkResult::advance(); - - llvm::StringMap renameMap; - for (func::FuncOp func : child.getOps()) { - if (!func.isPublic() || func.isDeclaration() || - func->hasAttr("pto.aicore")) - continue; - llvm::StringRef oldName = func.getSymName(); - std::string newName = (oldName + getPublicABISuffix(*kernelKind)).str(); - renameMap[oldName] = newName; - } - - for (func::FuncOp func : child.getOps()) { - func.walk([&](func::CallOp call) { - auto it = renameMap.find(call.getCallee()); - if (it != renameMap.end()) - call.setCallee(it->second); - }); - } - - for (func::FuncOp func : child.getOps()) { - auto it = renameMap.find(func.getSymName()); - if (it != renameMap.end()) - func.setSymName(it->second); - } - return WalkResult::advance(); - }); - return status; -} - -static LogicalResult verifyCompilationPlan(ModuleOp module, - bool cliBackendSpecified, - CompilationPlan &plan) { - if (failed(validatePTOBackendAttrs(module))) - return failure(); - if (failed(verifyBackendChildShape(module, cliBackendSpecified))) - return failure(); - - bool checkedVPTOChild = false; - for (ModuleOp child : module.getOps()) { - std::optional childBackend; - if (failed(parsePTOBackendAttr(child.getOperation(), childBackend))) - return failure(); - if (childBackend == PTOBackend::VPTO) { - checkedVPTOChild = true; - if (failed(verifyVPTOBackendChild(child))) - return failure(); - } else if (!childBackend && plan.effectiveBackend == PTOBackend::VPTO && - hasFunctionKernelKind(child)) { - checkedVPTOChild = true; - if (failed(verifyVPTOSourceSymbolNames(child))) - return failure(); - } - } - - if (!checkedVPTOChild && plan.effectiveBackend == PTOBackend::VPTO) - return verifyVPTOSourceSymbolNames(module); - - if (failed(verifyUniqueBackendABIExports(plan.children))) - return failure(); - if (failed(verifyExternalImportShape(plan.children))) - return failure(); - if (failed(resolveExternalImports(plan.children))) - return failure(); - - return success(); -} - -static LogicalResult buildCompilationPlan(ModuleOp module, - bool cliBackendSpecified, - PTOBackend cliOrDefaultBackend, - CompilationPlan &plan) { - plan = CompilationPlan{}; - plan.effectiveBackend = cliOrDefaultBackend; - - bool mixedBackendMode = false; - if (failed(resolveModuleBackend(module, cliBackendSpecified, mixedBackendMode, - plan.effectiveBackend))) - return failure(); - if (failed(collectBackendChildren(module, plan.children))) - return failure(); - plan.mode = mixedBackendMode ? CompilationMode::MixedFatobj - : getCompilationMode(plan.effectiveBackend, - plan.children); - if (failed(verifyCompilationPlan(module, cliBackendSpecified, plan))) - return failure(); - return success(); -} - -static LogicalResult emitSharedPreBackendSeamIR(ModuleOp module, - llvm::StringRef outputPath) { - if (outputPath.empty()) - return success(); - - if (outputPath == "-") { - module->print(llvm::outs()); - llvm::outs() << "\n"; - llvm::outs().flush(); - return success(); - } - - std::error_code ec; - llvm::ToolOutputFile outputFile(outputPath, ec, llvm::sys::fs::OF_None); - if (ec) { - llvm::errs() << "Error: failed to open seam IR file '" << outputPath - << "': " << ec.message() << "\n"; - return failure(); - } - - module->print(outputFile.os()); - outputFile.os() << "\n"; - outputFile.keep(); - return success(); -} - -static bool hasUnexpandedTileOps(ModuleOp module) { - bool found = false; - module.walk([&](Operation *op) { - if (found) - return; - if (isa(op)) - found = true; - }); - return found; -} - -static bool hasTilelangInlineHelpers(ModuleOp module) { - bool found = false; - module.walk([&](func::FuncOp func) { - if (found) - return; - if (func->hasAttr("pto.tilelang.inline_proc")) - found = true; - }); - return found; -} - -// -------------------------------------------------------------------------- -// Post-process C++ output: rewrite marker calls into Tile member calls. -// -// We emit marker calls in EmitC IR because EmitC currently does not provide a -// first-class op for member-function invocation. After translation, we rewrite: -// PTOAS__TILE_SET_VALUE(dst, offset, val) -> dst.SetValue(offset, val) -// PTOAS__TILE_GET_VALUE(src, offset) -> src.GetValue(offset) -// PTOAS__TILE_DATA(obj) -> obj.data() -// PTOAS__TILE_SET_VALIDSHAPE(obj, r, c) -> obj.SetValidShape(r, c) -// PTOAS__PTR_LOAD(ptr, offset) -> ptr[offset] -// PTOAS__PTR_STORE(ptr, offset, val) -> ptr[offset] = val -// PTOAS__EVENTID_ARRAY_LOAD(arr, idx) -> arr[idx] -// PTOAS__EVENTID_ARRAY_STORE(arr, idx, v) -> arr[idx] = v -// -------------------------------------------------------------------------- -struct ParsedMarkerCall { - size_t markerPos = std::string::npos; - size_t rparenPos = std::string::npos; - llvm::SmallVector args; -}; - -struct MarkerRewriteSpec { - llvm::StringRef marker; - llvm::StringRef memberName; - unsigned expectedNumArgs = 0; -}; - -struct MarkerSubscriptRewriteSpec { - llvm::StringRef marker; - unsigned expectedNumArgs = 0; - bool isStore = false; -}; - -static bool parseMarkerArgs(llvm::StringRef argsRef, - llvm::SmallVectorImpl &args) { - size_t partBegin = 0; - int parenDepth = 0; - for (size_t i = 0; i < argsRef.size(); ++i) { - char c = argsRef[i]; - if (c == '(') { - ++parenDepth; - continue; - } - if (c == ')') { - if (parenDepth > 0) - --parenDepth; - continue; - } - if (c == ',' && parenDepth == 0) { - args.push_back(argsRef.slice(partBegin, i).trim()); - partBegin = i + 1; - } - } - if (partBegin > argsRef.size()) - return false; - args.push_back(argsRef.drop_front(partBegin).trim()); - return true; -} - -static std::optional -findNextMarkerCall(const std::string &cpp, llvm::StringRef marker, - size_t searchPos) { - ParsedMarkerCall call; - call.markerPos = cpp.find(marker.str(), searchPos); - if (call.markerPos == std::string::npos) - return std::nullopt; - - size_t lparenPos = call.markerPos + marker.size(); - if (lparenPos >= cpp.size() || cpp[lparenPos] != '(') - return ParsedMarkerCall{call.markerPos, std::string::npos, {}}; - - size_t argsBegin = lparenPos + 1; - int parenDepth = 0; - for (size_t i = argsBegin; i < cpp.size(); ++i) { - char c = cpp[i]; - if (c == '(') { - ++parenDepth; - continue; - } - if (c != ')') - continue; - if (parenDepth == 0) { - call.rparenPos = i; - break; - } - --parenDepth; - } - if (call.rparenPos == std::string::npos) - return call; - - llvm::StringRef argsRef(cpp.data() + argsBegin, call.rparenPos - argsBegin); - if (!parseMarkerArgs(argsRef, call.args)) - call.args.clear(); - return call; -} - -template -static bool rewriteMarkerCalls(std::string &cpp, llvm::StringRef marker, - BuildReplacementFn buildReplacement) { - size_t searchPos = 0; - bool changed = false; - for (auto call = findNextMarkerCall(cpp, marker, searchPos); call; - call = findNextMarkerCall(cpp, marker, searchPos)) { - if (call->rparenPos == std::string::npos) { - searchPos = call->markerPos + marker.size(); - continue; - } - - std::optional replacement = buildReplacement(*call); - if (!replacement) { - searchPos = call->rparenPos + 1; - continue; - } - - cpp.replace(call->markerPos, (call->rparenPos - call->markerPos) + 1, - *replacement); - changed = true; - searchPos = call->markerPos + replacement->size(); - } - return changed; -} - -static bool rewriteMarkerCallToMember(std::string &cpp, llvm::StringRef marker, - llvm::StringRef memberName, - unsigned expectedNumArgs) { - return rewriteMarkerCalls( - cpp, marker, [&](const ParsedMarkerCall &call) -> std::optional { - if (call.args.size() != expectedNumArgs) - return std::nullopt; - - std::string replacement; - replacement.reserve(marker.size() + 16); - replacement.append(call.args[0].str()); - replacement.push_back('.'); - replacement.append(memberName.str()); - replacement.push_back('('); - if (expectedNumArgs >= 2) - replacement.append(call.args[1].str()); - if (expectedNumArgs == 3) { - replacement.append(", "); - replacement.append(call.args[2].str()); - } - replacement.push_back(')'); - return replacement; - }); -} - -static void rewriteMarkerCallsToMembers( - std::string &cpp, llvm::ArrayRef rewrites) { - bool changed = true; - while (changed) { - changed = false; - for (const MarkerRewriteSpec &rewrite : rewrites) { - changed |= rewriteMarkerCallToMember(cpp, rewrite.marker, - rewrite.memberName, - rewrite.expectedNumArgs); - } - } -} - -static void rewriteTileGetSetValueMarkers(std::string &cpp) { - static const MarkerRewriteSpec kTileMarkerRewrites[] = { - {"PTOAS__TILE_SET_VALUE", "SetValue", 3}, - {"PTOAS__TILE_GET_VALUE", "GetValue", 2}, - {"PTOAS__TILE_DATA", "data", 1}, - {"PTOAS__TILE_SET_VALIDSHAPE", "SetValidShape", 3}, - }; - rewriteMarkerCallsToMembers(cpp, kTileMarkerRewrites); -} - -static void rewriteAsyncEventMarkers(std::string &cpp) { - static const MarkerRewriteSpec kAsyncEventMarkerRewrites[] = { - {"PTOAS__ASYNC_EVENT_WAIT", "Wait", 2}, - {"PTOAS__ASYNC_EVENT_TEST", "Test", 2}, - }; - rewriteMarkerCallsToMembers(cpp, kAsyncEventMarkerRewrites); -} - -// -------------------------------------------------------------------------- -// EmitC cleanup: drop empty emitc.expression ops. -// -// After FormExpressions + CSE, EmitC expressions can become empty when their -// root op is CSE'd with an equivalent dominating value outside the expression -// region. Such expressions crash mlir::emitc::translateToCpp because -// ExpressionOp::getRootOp() returns nullptr. -// -------------------------------------------------------------------------- -static void dropEmptyEmitCExpressions(Operation *rootOp) { - llvm::SmallVector toErase; - rootOp->walk([&](emitc::ExpressionOp expr) { - if (expr.getRootOp()) - return; - Block *body = expr.getBody(); - if (!body) - return; - auto yield = dyn_cast(body->getTerminator()); - if (!yield || yield.getNumOperands() != 1) - return; - Value yielded = yield.getOperand(0); - expr.getResult().replaceAllUsesWith(yielded); - toErase.push_back(expr); - }); - for (emitc::ExpressionOp expr : llvm::reverse(toErase)) - expr.erase(); -} - -static Attribute getDefaultEmitCVariableInitAttr(OpBuilder &builder, Type type) { - if (auto intTy = dyn_cast(type)) - return builder.getIntegerAttr(intTy, 0); - if (isa(type)) - return builder.getIndexAttr(0); - if (auto floatTy = dyn_cast(type)) - return builder.getFloatAttr(floatTy, 0.0); - if (isa(type)) - return emitc::OpaqueAttr::get(builder.getContext(), ""); - return Attribute{}; -} - -// FormExpressions may inline conditions into emitc.expression, but the C++ -// emitter prints cf.br/cf.cond_br operands by variable name rather than by -// recursively emitting an expression. Materialize such operands so CFG-based -// lowering (e.g. scf.while -> cf.*) stays valid. -static void materializeControlFlowOperands(Operation *rootOp) { - llvm::SmallVector branches; - rootOp->walk([&](Operation *op) { - if (isa(op)) - branches.push_back(op); - }); - - OpBuilder builder(rootOp->getContext()); - for (Operation *op : branches) { - builder.setInsertionPoint(op); - for (OpOperand &operand : op->getOpOperands()) { - Value value = operand.get(); - auto expr = dyn_cast_or_null(value.getDefiningOp()); - if (!expr) - continue; - - Attribute initAttr = - getDefaultEmitCVariableInitAttr(builder, value.getType()); - if (!initAttr) - continue; - - Value tmp = - builder.create(op->getLoc(), value.getType(), - initAttr) - .getResult(); - builder.create(op->getLoc(), tmp, value); - operand.set(tmp); - } - } -} - -static bool rewriteMarkerCallToSubscript(std::string &cpp, llvm::StringRef marker, - unsigned expectedNumArgs, - bool isStore) { - return rewriteMarkerCalls( - cpp, marker, [&](const ParsedMarkerCall &call) -> std::optional { - if (call.args.size() != expectedNumArgs) - return std::nullopt; - if (isStore) { - return (call.args[0] + "[" + call.args[1] + "] = " + call.args[2]) - .str(); - } - return (call.args[0] + "[" + call.args[1] + "]").str(); - }); -} - -static void rewriteMarkerCallsToSubscripts( - std::string &cpp, llvm::ArrayRef rewrites) { - bool changed = true; - while (changed) { - changed = false; - for (const MarkerSubscriptRewriteSpec &rewrite : rewrites) { - changed |= rewriteMarkerCallToSubscript(cpp, rewrite.marker, - rewrite.expectedNumArgs, - rewrite.isStore); - } - } -} - -static void rewritePtrScalarMarkers(std::string &cpp) { - static const MarkerSubscriptRewriteSpec kPtrMarkerRewrites[] = { - {"PTOAS__PTR_LOAD", 2, false}, - {"PTOAS__PTR_STORE", 3, true}, - }; - rewriteMarkerCallsToSubscripts(cpp, kPtrMarkerRewrites); -} - -static void rewriteEventIdArrayMarkers(std::string &cpp) { - static const MarkerSubscriptRewriteSpec kEventIdMarkerRewrites[] = { - {"PTOAS__EVENTID_ARRAY_LOAD", 2, false}, - {"PTOAS__EVENTID_ARRAY_STORE", 3, true}, - }; - rewriteMarkerCallsToSubscripts(cpp, kEventIdMarkerRewrites); -} - -static bool rewriteAddPtrTraceMarkers(std::string &cpp, bool showTrace) { - size_t searchPos = 0; - bool changed = false; - for (auto call = findNextMarkerCall(cpp, "PTOAS__ADDPTR_TRACE", searchPos); - call; call = findNextMarkerCall(cpp, "PTOAS__ADDPTR_TRACE", searchPos)) { - if (call->rparenPos == std::string::npos) { - searchPos = call->markerPos + 1; - continue; - } - if (call->args.size() != 3) { - searchPos = call->rparenPos + 1; - continue; - } - - std::string replacement; - if (showTrace) { - replacement.reserve(64); - replacement.append("/* ADDPTR_TRACE: "); - replacement.append(call->args[0].str()); - replacement.append(" = "); - replacement.append(call->args[1].str()); - replacement.append(" + "); - replacement.append(call->args[2].str()); - replacement.append(" */"); - } - - size_t replaceEnd = call->rparenPos; - if (!showTrace) { - size_t i = call->rparenPos + 1; - while (i < cpp.size() && std::isspace(static_cast(cpp[i]))) - ++i; - if (i < cpp.size() && cpp[i] == ';') - replaceEnd = i; - } - - cpp.replace(call->markerPos, (replaceEnd - call->markerPos) + 1, - replacement); - changed = true; - searchPos = call->markerPos + replacement.size(); - } - return changed; -} - -static bool isGeneratedGlobalTensorDecl(llvm::StringRef trimmed, - llvm::StringRef &decl, - llvm::StringRef &varName) { - if (!trimmed.starts_with("GlobalTensor<") || !trimmed.ends_with(";") || - trimmed.contains('=') || trimmed.contains('(')) { - return false; - } - - decl = trimmed.drop_back().rtrim(); - size_t lastWs = decl.find_last_of(" \t"); - if (lastWs == llvm::StringRef::npos) - return false; - varName = decl.drop_front(lastWs + 1); - if (!varName.starts_with("v") || varName.size() <= 1) - return false; - return llvm::all_of(varName.drop_front(1), - [](char c) { return std::isdigit(c); }); -} - -static void rewriteHoistedGlobalTensorDecls(std::string &cpp) { - // When `declareVariablesAtTop` is enabled, the C++ emitter hoists SSA value - // declarations to the top of the function and emits assignments later. This - // requires the C++ type to be default-constructible. - // - // `GlobalTensor<...>` from pto-isa does NOT have a default constructor, so a - // hoisted declaration like: - // GlobalTensor<...> v42; - // fails to compile. Initialize those hoisted temporaries with a null pointer - // so they are constructible: - // GlobalTensor<...> v42(nullptr); - // - // We keep the assignment later; the null-initialized value is never used. - std::string out; - out.reserve(cpp.size() + 64); - - llvm::StringRef ref(cpp); - while (!ref.empty()) { - auto split = ref.split('\n'); - llvm::StringRef line = split.first; - llvm::StringRef rest = split.second; - - llvm::StringRef trimmed = line.trim(); - bool rewritten = false; - llvm::StringRef decl; - llvm::StringRef varName; - if (isGeneratedGlobalTensorDecl(trimmed, decl, varName)) { - size_t indentLen = line.find_first_not_of(" \t"); - if (indentLen == std::string::npos) - indentLen = 0; - llvm::StringRef indent = line.take_front(indentLen); - - out.append(indent.str()); - out.append(decl.str()); - out.append("(nullptr);"); - rewritten = true; - } - - if (!rewritten) - out.append(line.str()); - if (!rest.empty()) - out.push_back('\n'); - ref = rest; - } - - cpp.swap(out); -} - -namespace { -struct ConstantDeclCandidate { - size_t declLine = 0; - std::string indent; - std::string type; - bool hasInitializer = false; - std::string initializer; - size_t assignmentCount = 0; - size_t assignmentLine = 0; - std::string assignmentRhs; -}; -} // namespace - -static bool isGeneratedValueName(llvm::StringRef name) { - if (!name.consume_front("v") || name.empty()) - return false; - return llvm::all_of(name, [](char c) { return std::isdigit(c); }); -} - -static bool isConstFoldableScalarType(llvm::StringRef type) { - type = type.trim(); - if (type.starts_with("const ") || type.starts_with("constexpr ")) - return false; - return llvm::StringSwitch(type) - .Cases("bool", "float", "double", "half", "bfloat16_t", true) - .Cases("int8_t", "uint8_t", "int16_t", "uint16_t", true) - .Cases("int32_t", "uint32_t", "int64_t", "uint64_t", true) - .Default(false); -} - -static bool isLiteralInitializer(llvm::StringRef rhs) { - rhs = rhs.trim(); - if (rhs.empty()) - return false; - if (rhs == "true" || rhs == "false" || rhs == "nullptr") - return true; - - static const llvm::Regex kIntLiteral( - R"(^[+-]?(0[xX][0-9A-Fa-f]+|[0-9]+)[uUlL]*$)"); - static const llvm::Regex kFloatLiteral( - R"(^[+-]?(([0-9]+\.[0-9]*|\.[0-9]+|[0-9]+)([eE][+-]?[0-9]+)?|[0-9]+[eE][+-]?[0-9]+)[fF]?$)"); - static const llvm::Regex kHexFloatLiteral( - R"(^[+-]?0[xX]([0-9A-Fa-f]+\.[0-9A-Fa-f]*|[0-9A-Fa-f]+|\.[0-9A-Fa-f]+)[pP][+-]?[0-9]+[fF]?$)"); - static const llvm::Regex kSpecialFloatLiteral( - R"(^[+-]?(nan|inf)[fF]?$)"); - - return kIntLiteral.match(rhs) || kFloatLiteral.match(rhs) || - kHexFloatLiteral.match(rhs) || kSpecialFloatLiteral.match(rhs); -} - -static std::string normalizeConstInitializer(llvm::StringRef type, - llvm::StringRef rhs) { - type = type.trim(); - rhs = rhs.trim(); - if (type == "bool") { - if (rhs == "0" || rhs == "false") - return "false"; - if (rhs == "1" || rhs == "-1" || rhs == "true") - return "true"; - } - return rhs.str(); -} - -static bool parseConstantDeclarationLine(llvm::StringRef line, - ConstantDeclCandidate &candidate, - std::string &valueName) { - llvm::StringRef trimmed = line.trim(); - if (trimmed.empty() || trimmed.starts_with("#") || trimmed.starts_with("//") || - !trimmed.ends_with(";")) - return false; - - llvm::StringRef body = trimmed.drop_back().rtrim(); - if (body.starts_with("return") || body.starts_with("goto ") || - body.starts_with("if ") || body.starts_with("if(") || - body.starts_with("switch ") || body.starts_with("switch(") || - body.starts_with("for ") || body.starts_with("for(") || - body.starts_with("while ") || body.starts_with("while(") || - body.starts_with("case ") || body == "default") - return false; - - llvm::StringRef lhs = body; - llvm::StringRef rhs; - if (size_t eqPos = body.find('='); eqPos != llvm::StringRef::npos) { - lhs = body.take_front(eqPos).rtrim(); - rhs = body.drop_front(eqPos + 1).trim(); - } - - size_t lastWs = lhs.find_last_of(" \t"); - if (lastWs == llvm::StringRef::npos) - return false; - - llvm::StringRef type = lhs.take_front(lastWs).rtrim(); - llvm::StringRef name = lhs.drop_front(lastWs + 1).trim(); - if (!isGeneratedValueName(name) || !isConstFoldableScalarType(type)) - return false; - - size_t indentLen = line.find_first_not_of(" \t"); - if (indentLen == llvm::StringRef::npos) - indentLen = 0; - candidate.indent = line.take_front(indentLen).str(); - candidate.type = type.str(); - valueName = name.str(); - - if (!rhs.empty()) { - if (!isLiteralInitializer(rhs)) - return false; - candidate.hasInitializer = true; - candidate.initializer = normalizeConstInitializer(type, rhs); - } - - return true; -} - -static bool parseGeneratedValueAssignment(llvm::StringRef line, - llvm::StringRef &valueName, - llvm::StringRef &rhs) { - llvm::StringRef trimmed = line.trim(); - if (trimmed.empty() || trimmed.starts_with("#") || trimmed.starts_with("//") || - !trimmed.ends_with(";")) - return false; - - llvm::StringRef body = trimmed.drop_back().rtrim(); - size_t eqPos = body.find('='); - if (eqPos == llvm::StringRef::npos) - return false; - - llvm::StringRef lhs = body.take_front(eqPos).rtrim(); - rhs = body.drop_front(eqPos + 1).trim(); - if (!isGeneratedValueName(lhs)) - return false; - valueName = lhs; - return true; -} - -static void rewriteScalarConstantDecls(std::string &cpp) { - llvm::SmallVector lines; - for (llvm::StringRef ref(cpp); !ref.empty(); ref = ref.split('\n').second) { - auto split = ref.split('\n'); - lines.push_back(split.first.str()); - } - - llvm::SmallVector eraseLine(lines.size(), false); - auto rewriteSegment = [&](size_t beginLine, size_t endLine) { - llvm::StringMap candidates; - - for (size_t i = beginLine; i <= endLine; ++i) { - ConstantDeclCandidate candidate; - std::string valueName; - if (parseConstantDeclarationLine(lines[i], candidate, valueName)) { - candidate.declLine = i; - candidates[valueName] = std::move(candidate); - continue; - } - - llvm::StringRef assignedName; - llvm::StringRef rhs; - if (!parseGeneratedValueAssignment(lines[i], assignedName, rhs)) - continue; - - auto it = candidates.find(assignedName); - if (it == candidates.end()) - continue; - - ConstantDeclCandidate &info = it->second; - ++info.assignmentCount; - info.assignmentLine = i; - info.assignmentRhs = rhs.str(); - } - - for (auto &entry : candidates) { - llvm::StringRef valueName = entry.getKey(); - ConstantDeclCandidate &info = entry.getValue(); - - std::string initializer; - if (info.hasInitializer) { - if (info.assignmentCount != 0) - continue; - initializer = info.initializer; - } else { - if (info.assignmentCount != 1) - continue; - if (!isLiteralInitializer(info.assignmentRhs)) - continue; - initializer = normalizeConstInitializer( - info.type, llvm::StringRef(info.assignmentRhs)); - eraseLine[info.assignmentLine] = true; - } - - lines[info.declLine] = (info.indent + "const " + info.type + " " + - valueName.str() + " = " + initializer + ";"); - } - }; - - int braceDepth = 0; - size_t segmentStart = 0; - for (size_t i = 0; i < lines.size(); ++i) { - int depthBefore = braceDepth; - for (char c : lines[i]) { - if (c == '{') - ++braceDepth; - else if (c == '}') - --braceDepth; - } - - if (depthBefore == 0 && braceDepth > 0) - segmentStart = i; - if (depthBefore > 0 && braceDepth == 0) - rewriteSegment(segmentStart, i); - } - - std::string out; - out.reserve(cpp.size()); - for (size_t i = 0; i < lines.size(); ++i) { - if (eraseLine[i]) - continue; - out.append(lines[i]); - if (i + 1 != lines.size()) - out.push_back('\n'); - } - cpp.swap(out); -} - -static bool shouldDeclareVariablesAtTop(ModuleOp module) { - auto hasMultiBlockFunc = [](auto func) { return func.getBlocks().size() > 1; }; - return llvm::any_of(module.getOps(), hasMultiBlockFunc) || - llvm::any_of(module.getOps(), hasMultiBlockFunc); -} - -static LogicalResult emitEmitCSourceArtifact(ModuleOp module, - bool emitAddPtrTraceComments, - std::string &cppOutput, - llvm::raw_ostream &diagOS) { - dropEmptyEmitCExpressions(module); - materializeControlFlowOperands(module); - if (failed(reorderEmitCFunctions(module))) { - diagOS << "Error: Failed to order emitted functions for C++ emission.\n"; - return failure(); - } - - cppOutput.clear(); - llvm::raw_string_ostream cppOS(cppOutput); - bool declareVariablesAtTop = shouldDeclareVariablesAtTop(module); - if (failed(emitc::translateToCpp( - module, cppOS, /*declareVariablesAtTop=*/declareVariablesAtTop))) { - diagOS << "Error: Failed to emit C++.\n"; - return failure(); - } - cppOS.flush(); - rewriteTileGetSetValueMarkers(cppOutput); - rewriteAsyncEventMarkers(cppOutput); - rewritePtrScalarMarkers(cppOutput); - rewriteEventIdArrayMarkers(cppOutput); - rewriteAddPtrTraceMarkers(cppOutput, emitAddPtrTraceComments); - rewriteScalarConstantDecls(cppOutput); - rewriteHoistedGlobalTensorDecls(cppOutput); - return success(); -} - -static LogicalResult configureSharedPTOPipeline(PassManager &pm, - PTOBuildLevel effectiveLevel) { - if (failed(applyPassManagerCLOptions(pm))) - return failure(); - - pm.addNestedPass( - pto::createPTOAssignDefaultFrontendPipeIdPass()); - pm.addNestedPass( - pto::createPTOLowerFrontendPipeOpsPass()); - // pm.addNestedPass(pto::createPTOVerifyTFreePass()); - pm.addPass(pto::createPTOInferValidatePipeInitPass()); - pm.addNestedPass(pto::createLoweringSyncToPipePass()); - - if (!disableInferLayout) - pm.addNestedPass(pto::createInferPTOLayoutPass()); - pm.addNestedPass(pto::createPTOA5NormalizeTMovPass()); - pm.addPass(pto::createPTOViewToMemrefPass()); - - if (effectiveLevel != PTOBuildLevel::Level3) { - PlanMemoryOptions planMemoryOption; - planMemoryOption.memMode = MemPlanMode::LOCAL_MEM_PLAN; - planMemoryOption.enableGlobalReuse = false; - planMemoryOption.enablePrintMemoryAllocatedSize = false; - pm.addPass(pto::createPlanMemoryPass(planMemoryOption)); - } - pm.addPass(pto::createPTOResolveReservedBuffersPass()); - - // Conditionally add one automatic synchronization mode. Barrier-all is a - // conservative standalone pass; InsertSync and GraphSyncSolver are set/wait - // solvers. - if (enableInsertSync) - pm.addNestedPass(pto::createPTOInsertSyncPass()); - else if (enableInjectBarrierAllSync) - pm.addNestedPass( - pto::createPTOInjectBarrierAllSyncPass()); - else if (enableGraphSyncSolver) { - PTOGraphSyncSolverOptions graphSyncOpts; - graphSyncOpts.eventIdNumMax = graphSyncSolverEventIdMax; - pm.addNestedPass( - pto::createPTOGraphSyncSolverPass(graphSyncOpts)); - } - - return success(); -} - -static void addSharedPTOSeamPasses(PassManager &pm) { - // Reintroduce tile-native handles once on the shared mainline so both - // backends consume the same post-planning seam IR. - pm.addPass(pto::createPTOMaterializeTileHandlesPass()); - pm.addPass(createCSEPass()); -} - -static void prepareVPTOForEmission(PassManager &pm) { - auto &kernelModulePM = pm.nest(); - kernelModulePM.addPass(createCanonicalizerPass()); - kernelModulePM.addPass(createCSEPass()); - kernelModulePM.addPass(pto::createVPTOPtrNormalizePass()); - kernelModulePM.addPass(pto::createVPTOPtrCastCleanupPass()); - kernelModulePM.addPass(createReconcileUnrealizedCastsPass()); - kernelModulePM.addNestedPass( - createVPTOExpandWrapperOpsPass()); - kernelModulePM.addPass(createCSEPass()); - kernelModulePM.addNestedPass( - pto::createPTOInferVPTOVecScopePass()); - kernelModulePM.addPass(createCanonicalizerPass()); - kernelModulePM.addPass(createCSEPass()); - kernelModulePM.addPass(pto::createPTOValidateVPTOEmissionIRPass()); -} - -static void lowerPTOToVPTOBackend(PassManager &pm, int argc, char **argv) { - // TileOp Expand path: - // 1. ExpandTileOp: instantiate TileLang DSL templates, replace tile ops - // with func.call to template functions (tile_buf params) - // 2. InlineLibCall: inline template function bodies - // 3. FoldTileBufIntrinsics: fold tile_buf_addr / tile_valid_rows / - // tile_valid_cols to concrete memref/constant values - auto &kernelModulePM = pm.nest(); - pto::ExpandTileOpOptions expandOpts = resolveExpandTileOpOptions(argc, argv); - kernelModulePM.addPass(pto::createExpandTileOpPass(expandOpts)); - - kernelModulePM.addPass(pto::createPTOInlineLibCallPass()); - kernelModulePM.addNestedPass( - pto::createFoldTileBufIntrinsicsPass()); - // FoldTileBufIntrinsics materializes many constant branch conditions. - // Clean them up immediately on the TileOp expansion path before the - // authoring-stage VPTO verifier and let the existing CSE passes remove the - // resulting dead values later in the pipeline. - kernelModulePM.addPass(mlir::createSCCPPass()); - kernelModulePM.addPass(mlir::createCanonicalizerPass()); -} - -static void inlineTilelangHelpersOnVPTOInput(PassManager &pm) { - auto &kernelModulePM = pm.nest(); - kernelModulePM.addPass(pto::createPTOInlineLibCallPass()); - kernelModulePM.addPass(mlir::createSCCPPass()); - kernelModulePM.addPass(mlir::createCanonicalizerPass()); -} - -static pto::VPTOEmissionOptions buildVPTOEmissionOptions() { - pto::VPTOEmissionOptions options; - options.dumpVPTOIR = false; - options.targetTriple = "hiipu64-hisilicon-cce"; - return options; -} - -struct CompiledBackendChild { - const BackendChildPlan *plan = nullptr; - PTOBackend backend = PTOBackend::EmitC; - std::optional kernelKind; - ModuleOp vptoModule; - std::string cppSource; - pto::EmittedLLVMModule cubeModule; - pto::EmittedLLVMModule vectorModule; -}; - -static LogicalResult runVPTOBackendPipeline(ModuleOp module, int argc, - char **argv, - bool hasTileOpsToExpand, - bool hasTilelangHelpers); - -static LogicalResult runSharedPTOToSeamPipeline(ModuleOp module, - PTOBuildLevel effectiveLevel, - llvm::StringRef arch, - llvm::raw_ostream &diagOS) { - PassManager pm(module.getContext()); - if (failed(configureSharedPTOPipeline(pm, effectiveLevel))) - return failure(); - addSharedPTOSeamPasses(pm); - if (failed(applyConfiguredPassManagerCLOptions(pm, "mixed child PTOAS pipeline"))) - return failure(); - module->setAttr("pto.target_arch", StringAttr::get(module.getContext(), arch)); - if (failed(pm.run(module))) { - diagOS << "Error: mixed child PTO pass execution failed.\n"; - return failure(); - } - return success(); -} - -static LogicalResult reconcileExternalDeclarationTypes(ModuleOp module) { - SymbolTable symbolTable(module); - llvm::StringMap> operandTypesByExternalCallee; - LogicalResult status = success(); - - module.walk([&](func::CallOp call) { - if (failed(status)) - return WalkResult::interrupt(); - auto callee = symbolTable.lookup(call.getCallee()); - if (!callee || !callee.isDeclaration()) - return WalkResult::advance(); - - SmallVector operandTypes(call.getOperandTypes().begin(), - call.getOperandTypes().end()); - auto [it, inserted] = - operandTypesByExternalCallee.try_emplace(call.getCallee(), operandTypes); - if (!inserted && it->second != operandTypes) { - status = call.emitError() - << "external declaration '" << call.getCallee() - << "' is called with inconsistent operand types after PTO " - "normalization"; - return WalkResult::interrupt(); - } - return WalkResult::advance(); - }); - if (failed(status)) - return failure(); - - for (func::FuncOp func : module.getOps()) { - if (!func.isDeclaration()) - continue; - auto it = operandTypesByExternalCallee.find(func.getSymName()); - if (it == operandTypesByExternalCallee.end()) - continue; - if (func.getResultTypes().empty() && - llvm::equal(func.getArgumentTypes(), it->second)) - continue; - if (!func.getResultTypes().empty()) { - return func.emitError() - << "external declarations with results are not supported in mixed " - "backend object emission"; - } - func.setFunctionType( - FunctionType::get(func.getContext(), it->second, TypeRange{})); - } - return success(); -} - -static LogicalResult runEmitCBackendPipeline(ModuleOp module, - PTOBuildLevel effectiveLevel, - llvm::StringRef arch, - llvm::raw_ostream &diagOS) { - PassManager sharedPM(module.getContext()); - // Shared PTO normalization can rewrite pto.aicore argument types before the - // private external declarations that call them are reconciled below. - sharedPM.enableVerifier(false); - if (failed(configureSharedPTOPipeline(sharedPM, effectiveLevel))) - return failure(); - addSharedPTOSeamPasses(sharedPM); - if (failed(applyConfiguredPassManagerCLOptions(sharedPM, - "mixed child shared pipeline"))) - return failure(); - module->setAttr("pto.target_arch", StringAttr::get(module.getContext(), arch)); - if (failed(sharedPM.run(module))) { - diagOS << "Error: mixed EmitC child shared pass execution failed.\n"; - return failure(); - } - - if (failed(reconcileExternalDeclarationTypes(module))) - return failure(); - - PassManager emitCPM(module.getContext()); - if (arch == "a3") - emitCPM.addPass(pto::createEmitPTOManualPass(pto::PTOArch::A3)); - else - emitCPM.addPass(pto::createEmitPTOManualPass(pto::PTOArch::A5)); - emitCPM.addPass(emitc::createFormExpressionsPass()); - emitCPM.addPass(mlir::createCSEPass()); - if (failed(applyConfiguredPassManagerCLOptions(emitCPM, - "mixed child EmitC pipeline"))) - return failure(); - if (failed(emitCPM.run(module))) { - diagOS << "Error: mixed EmitC child pass execution failed.\n"; - return failure(); - } - return success(); -} - -static LogicalResult compileMixedEmitCChild(const BackendChildPlan &plan, - PTOBuildLevel effectiveLevel, - llvm::StringRef arch, - CompiledBackendChild &compiled, - llvm::raw_ostream &diagOS) { - ModuleOp module = plan.module; - if (failed(runEmitCBackendPipeline(module, effectiveLevel, arch, diagOS))) - return failure(); - - compiled.plan = &plan; - compiled.backend = PTOBackend::EmitC; - compiled.kernelKind = plan.kernelKind; - if (failed(emitEmitCSourceArtifact(module, emitAddPtrTrace, - compiled.cppSource, diagOS))) - return failure(); - return success(); -} - -static LogicalResult compileMixedVPTOChild(const BackendChildPlan &plan, - PTOBuildLevel effectiveLevel, - llvm::StringRef arch, int argc, - char **argv, - CompiledBackendChild &compiled, - llvm::raw_ostream &diagOS) { - ModuleOp module = plan.module; - const bool hasTileOpsToExpand = hasUnexpandedTileOps(module); - const bool hasTilelangHelpers = hasTilelangInlineHelpers(module); - - if (hasTileOpsToExpand) { - if (failed(runSharedPTOToSeamPipeline(module, effectiveLevel, arch, diagOS))) - return failure(); - } - - if (failed(runVPTOBackendPipeline(module, argc, argv, hasTileOpsToExpand, - hasTilelangHelpers))) - return failure(); - - compiled.plan = &plan; - compiled.backend = PTOBackend::VPTO; - compiled.kernelKind = plan.kernelKind; - compiled.vptoModule = module; - - pto::VPTOEmissionOptions options = buildVPTOEmissionOptions(); - if (failed(pto::lowerVPTOModuleToLLVMModules( - module, options, compiled.cubeModule, compiled.vectorModule, diagOS))) { - diagOS << "Error: Failed to lower mixed VPTO child to LLVM modules.\n"; - return failure(); - } - return success(); -} - -static LogicalResult compilePTOChildren(CompilationPlan &plan, - PTOBuildLevel effectiveLevel, - llvm::StringRef arch, int argc, - char **argv, - SmallVectorImpl - &compiledChildren, - llvm::raw_ostream &diagOS) { - for (BackendChildPlan &child : plan.children) { - if (failed(applyResolvedExternalImports(child))) - return failure(); - - CompiledBackendChild compiled; - if (child.backend == PTOBackend::EmitC) { - if (failed(compileMixedEmitCChild(child, effectiveLevel, arch, compiled, - diagOS))) - return failure(); - } else { - if (failed(compileMixedVPTOChild(child, effectiveLevel, arch, argc, argv, - compiled, diagOS))) - return failure(); - } - compiledChildren.push_back(std::move(compiled)); - } - return success(); -} - -static LogicalResult createTempPath(llvm::StringRef prefix, - llvm::StringRef suffix, - std::string &path, - llvm::raw_ostream &diagOS) { - llvm::SmallString<128> tempPath; - int fd = -1; - std::error_code ec = - llvm::sys::fs::createTemporaryFile(prefix, suffix, fd, tempPath); - if (ec) { - diagOS << "Error: failed to create temporary file for " << prefix << suffix - << ": " << ec.message() << "\n"; - return failure(); - } - llvm::sys::Process::SafelyCloseFileDescriptor(fd); - path = tempPath.str().str(); - return success(); -} - -static void removeTempPaths(ArrayRef paths) { - for (const std::string &path : paths) - if (!path.empty()) - llvm::sys::fs::remove(path); -} - -static LogicalResult compileMixedDeviceObjects( - ArrayRef compiledChildren, - const ObjectEmissionToolchain &toolchain, - SmallVectorImpl &artifacts, - SmallVectorImpl &tempPaths, llvm::raw_ostream &diagOS) { - for (const CompiledBackendChild &child : compiledChildren) { - if (child.backend == PTOBackend::EmitC) { - std::string cppPath; - std::string objPath; - std::string stderrPath; - if (failed(createTempPath("ptoas-mixed-emitc", ".cpp", cppPath, diagOS)) || - failed(createTempPath("ptoas-mixed-emitc", ".o", objPath, diagOS)) || - failed(createTempPath("ptoas-mixed-emitc", ".log", stderrPath, - diagOS))) - return failure(); - tempPaths.push_back(cppPath); - tempPaths.push_back(objPath); - tempPaths.push_back(stderrPath); - - if (failed(emitCppFatobj(child.cppSource, cppPath, objPath, toolchain, - stderrPath, diagOS))) - return failure(); - - BackendArtifact artifact; - artifact.backend = PTOBackend::EmitC; - artifact.kind = BackendArtifactKind::Fatobj; - artifact.path = objPath; - if (child.plan) - for (const BackendExportPlan &exportPlan : child.plan->exports) - artifact.exportedABISymbols.push_back(exportPlan.abiName); - artifacts.push_back(std::move(artifact)); - continue; - } - - auto emitVPTOObject = [&](llvm::Module *llvmModule, - ObjectEmissionDeviceTarget target, - BackendArtifactKind artifactKind) -> LogicalResult { - if (!llvmModule) - return success(); - std::string llPath; - std::string objPath; - std::string stderrPath; - if (failed(createTempPath("ptoas-mixed-vpto", ".ll", llPath, diagOS)) || - failed(createTempPath("ptoas-mixed-vpto", ".o", objPath, diagOS)) || - failed(createTempPath("ptoas-mixed-vpto", ".log", stderrPath, diagOS))) - return failure(); - tempPaths.push_back(llPath); - tempPaths.push_back(objPath); - tempPaths.push_back(stderrPath); - - if (failed(writeLLVMModule(*llvmModule, llPath, diagOS))) - return failure(); - if (failed(compileLLVMToDeviceObject(llPath, objPath, target, toolchain, - stderrPath, diagOS))) - return failure(); - - BackendArtifact artifact; - artifact.backend = PTOBackend::VPTO; - artifact.kind = artifactKind; - artifact.path = objPath; - if (child.plan) - for (const BackendExportPlan &exportPlan : child.plan->exports) - artifact.exportedABISymbols.push_back(exportPlan.abiName); - artifacts.push_back(std::move(artifact)); - return success(); - }; - - if (failed(emitVPTOObject(child.vectorModule.module.get(), - ObjectEmissionDeviceTarget::Vector, - BackendArtifactKind::VectorObject))) - return failure(); - if (failed(emitVPTOObject(child.cubeModule.module.get(), - ObjectEmissionDeviceTarget::Cube, - BackendArtifactKind::CubeObject))) - return failure(); - } - - if (artifacts.empty()) { - diagOS << "Error: mixed fatobj compilation produced no device objects.\n"; - return failure(); - } - return success(); -} - -static LogicalResult buildMixedHostStubSource( - ArrayRef compiledChildren, std::string &stubSource, - llvm::raw_ostream &diagOS) { - SmallVector stubModules; - for (const CompiledBackendChild &child : compiledChildren) { - if (child.backend != PTOBackend::VPTO || !child.vptoModule) - continue; - bool hasAICoreKernel = false; - ModuleOp vptoModule = child.vptoModule; - for (func::FuncOp func : vptoModule.getOps()) { - if (!func.isExternal() && func->hasAttr("pto.aicore")) { - hasAICoreKernel = true; - break; - } - } - if (hasAICoreKernel) - stubModules.push_back(child.vptoModule); - } - if (stubModules.empty()) { - stubSource = "#ifndef __global__\n#define __global__\n#endif\n\n" - "#ifndef __gm__\n#define __gm__\n#endif\n\n"; - return success(); - } - if (failed(pto::emitVPTOHostStubSource(stubModules, stubSource, diagOS))) { - diagOS << "Error: Failed to emit mixed VPTO host stub source.\n"; - return failure(); - } - return success(); -} - -static LogicalResult linkMixedFatobj( - ArrayRef artifacts, llvm::StringRef stubSource, - llvm::StringRef outputPath, const ObjectEmissionToolchain &toolchain, - SmallVectorImpl &tempPaths, llvm::raw_ostream &diagOS) { - std::string stubPath; - std::string stderrPath; - if (failed(createTempPath("ptoas-mixed-stub", ".cpp", stubPath, diagOS)) || - failed(createTempPath("ptoas-mixed", ".log", stderrPath, diagOS))) - return failure(); - tempPaths.push_back(stubPath); - tempPaths.push_back(stderrPath); - - if (failed(writeHostStubSource(stubSource, stubPath, diagOS))) - return failure(); - - std::string sanitizedModuleId = - outputPath.empty() || outputPath == "-" ? "ptoas_mixed_fatobj" - : outputPath.str(); - for (char &c : sanitizedModuleId) - if (!std::isalnum(static_cast(c)) && c != '_') - c = '_'; - - SmallVector fatobjPaths; - for (size_t i = 0, e = artifacts.size(); i < e; ++i) { - const BackendArtifact &artifact = artifacts[i]; - if (artifact.kind == BackendArtifactKind::Fatobj) { - fatobjPaths.push_back(artifact.path); - continue; - } - - std::string mergedObjPath; - std::string fatobjPath; - if (failed(createTempPath("ptoas-mixed-backend-merged", ".o", mergedObjPath, - diagOS)) || - failed(createTempPath("ptoas-mixed-backend-fatobj", ".o", fatobjPath, - diagOS))) - return failure(); - tempPaths.push_back(mergedObjPath); - tempPaths.push_back(fatobjPath); - - if (failed(mergeDeviceObjects(ArrayRef(artifact.path), - mergedObjPath, toolchain, - stderrPath, diagOS))) - return failure(); - - std::string groupModuleId = sanitizedModuleId + "_backend_" + - std::to_string(i); - if (failed(compileStubToFatobj(stubPath, mergedObjPath, fatobjPath, - groupModuleId, toolchain, stderrPath, - diagOS))) - return failure(); - fatobjPaths.push_back(fatobjPath); - } - - if (fatobjPaths.empty()) { - diagOS << "Error: mixed fatobj linking produced no backend fatobjs.\n"; - return failure(); - } - if (fatobjPaths.size() == 1) { - if (std::error_code ec = - llvm::sys::fs::copy_file(fatobjPaths.front(), outputPath)) { - diagOS << "Error: failed to copy mixed backend fatobj to " << outputPath - << ": " << ec.message() << "\n"; - return failure(); - } - return success(); - } - - if (failed(linkFatobjs(fatobjPaths, outputPath, toolchain, stderrPath, - diagOS))) - return failure(); - return success(); -} - -static int runMixedFatobjPlan(CompilationPlan &plan, - PTOBuildLevel effectiveLevel, - llvm::StringRef arch, int argc, char **argv, - llvm::StringRef outputPath, - llvm::raw_ostream &diagOS) { - if (emitMlirIR || emitVPTO || ptoPrintSeamIR || !ptoSeamIRFile.empty()) { - diagOS << "Error: mixed pto.backend fatobj mode does not support " - "debug IR output flags.\n"; - return 1; - } - if (outputPath.empty() || outputPath == "-") { - diagOS << "Error: mixed pto.backend fatobj mode requires an explicit " - "file path passed with -o.\n"; - return 1; - } - - SmallVector compiledChildren; - if (failed(compilePTOChildren(plan, effectiveLevel, arch, argc, argv, - compiledChildren, diagOS))) - return 1; - - std::string stubSource; - if (failed(buildMixedHostStubSource(compiledChildren, stubSource, diagOS))) - return 1; - - ObjectEmissionToolchain toolchain; - if (failed(discoverObjectEmissionToolchain(toolchain, diagOS))) - return 1; - - SmallVector tempPaths; - auto cleanup = llvm::make_scope_exit([&]() { removeTempPaths(tempPaths); }); - - SmallVector artifacts; - if (failed(compileMixedDeviceObjects(compiledChildren, toolchain, artifacts, - tempPaths, diagOS))) - return 1; - - if (failed(linkMixedFatobj(artifacts, stubSource, outputPath, toolchain, - tempPaths, diagOS))) - return 1; - return 0; -} - -static int emitVPTOBackendResult(ModuleOp module, - llvm::ToolOutputFile &outputFile) { - if (emitVPTO) { - module.print(outputFile.os()); - outputFile.os() << "\n"; - outputFile.keep(); - return 0; - } - - pto::VPTOEmissionOptions options = buildVPTOEmissionOptions(); - std::string stubSource; - if (failed(pto::emitVPTOHostStubSource(module, stubSource, llvm::errs()))) { - llvm::errs() << "Error: Failed to emit VPTO host stub source.\n"; - return 1; - } - - pto::EmittedLLVMModule cubeModule; - pto::EmittedLLVMModule vectorModule; - if (failed( - pto::lowerVPTOModuleToLLVMModules(module, options, cubeModule, - vectorModule, llvm::errs()))) { - llvm::errs() << "Error: Failed to lower VPTO to LLVM modules.\n"; - return 1; - } - - if (failed(pto::emitVPTOFatobjFromLLVMModules( - cubeModule.module.get(), vectorModule.module.get(), stubSource, - outputFile, llvm::errs()))) { - llvm::errs() << "Error: Failed to emit VPTO fatobj.\n"; - return 1; - } - outputFile.keep(); - return 0; -} - -static LogicalResult runVPTOBackendPipeline(ModuleOp module, int argc, - char **argv, - bool hasTileOpsToExpand, - bool hasTilelangHelpers) { - PassManager pm(module.getContext()); - pm.enableVerifier(); - pm.addPass(pto::createVPTOSplitCVModulePass()); - pm.addPass(pto::createVPTONormalizeContainerPass()); - if (!hasTileOpsToExpand && hasTilelangHelpers) - inlineTilelangHelpersOnVPTOInput(pm); - if (hasTileOpsToExpand) - lowerPTOToVPTOBackend(pm, argc, argv); - prepareVPTOForEmission(pm); - if (failed(applyConfiguredPassManagerCLOptions( - pm, "VPTO unified emission pipeline"))) - return failure(); - if (failed(pm.run(module))) { - llvm::errs() << "Error: VPTO emission pipeline failed.\n"; - return failure(); - } - if (failed(applyVPTOPublicABIExportNames(module))) - return failure(); - return success(); -} - -int mlir::pto::runPTOASDriver(int argc, char **argv) { - DialectRegistry registry; - registry.insert(); - registry.insert(); - registry.insert(); - registry.insert(); - registry.insert(); - registry.insert(); - registry.insert(); - registry.insert(); - registry.insert(); - - registry.insert(); - arith::registerBufferizableOpInterfaceExternalModels(registry); - tensor::registerBufferizableOpInterfaceExternalModels(registry); - pto::registerBufferizableOpInterfaceExternalModels(registry); - - registry.insert(); - registry.insert(); - mlir::registerAllPasses(); - ::registerPTOPasses(); - mlir::pto::registerPTOViewToMemrefPass(); - ::registerPTOInlineLibCall(); - ::registerFoldTileBufIntrinsics(); - ::registerExpandTileOp(); - mlir::registerPassManagerCLOptions(); - - llvm::cl::SetVersionPrinter(printPTOASVersion); - - bool cliArchSpecified = false; - for (int i = 1; i < argc; ++i) { - llvm::StringRef arg(argv[i]); - if (arg == "--pto-arch" || arg.starts_with("--pto-arch=")) - cliArchSpecified = true; - } - bool cliBackendSpecified = hasCLIOption(argc, argv, "--pto-backend"); - - // Parse command line options - llvm::cl::ParseCommandLineOptions(argc, argv, "PTO Assembler (ptoas)\n"); - - PTOBackend effectiveBackend = PTOBackend::EmitC; - if (!parseBackend(ptoBackend, effectiveBackend)) { - llvm::errs() << "Error: invalid --pto-backend='" << ptoBackend - << "'. Expected 'emitc' or 'vpto'.\n"; - return 1; - } - - // Read whole input first (so we can auto-detect .ptobc by magic). - auto fileOrErr = llvm::MemoryBuffer::getFileOrSTDIN(inputFilename); - if (!fileOrErr) { - llvm::errs() << "Error: Could not open input file: " - << fileOrErr.getError().message() << "\n"; - return 1; - } - - MLIRContext context(registry); - // Be tolerant: ptobc decode may materialize ops from dialects that aren't - // explicitly registered/loaded in this tool yet. - context.allowUnregisteredDialects(true); - - context.getOrLoadDialect(); - context.getOrLoadDialect(); - context.getOrLoadDialect(); - context.getOrLoadDialect(); - context.getOrLoadDialect(); - context.getOrLoadDialect(); - context.getOrLoadDialect(); - context.getOrLoadDialect(); - - OwningOpRef module; - llvm::StringRef buf = (*fileOrErr)->getBuffer(); - const bool isPTOBC = (buf.size() >= 6 && std::memcmp(buf.data(), "PTOBC\0", 6) == 0); - - auto normalizeArch = [](llvm::StringRef archValue) { - std::string normalized = archValue.str(); - for (char &c : normalized) - c = static_cast(std::tolower(static_cast(c))); - return normalized; - }; - auto detectTextualModuleArch = [&](llvm::StringRef text) -> std::optional { - llvm::SmallVector matches; - llvm::Regex archRegex( - R"ptoarch("?(pto\.target_arch)"?[[:space:]]*=[[:space:]]*"([[:alpha:][:digit:]_]+)")ptoarch"); - if (!archRegex.match(text, &matches) || matches.size() < 3) - return std::nullopt; - return normalizeArch(matches[2]); - }; - - std::string arch = normalizeArch(ptoTargetArch); - if (cliArchSpecified) { - if (arch != "a3" && arch != "a5") { - llvm::errs() << "Error: invalid --pto-arch='" << ptoTargetArch - << "'. Expected 'a3' or 'a5'.\n"; - return 1; - } - } else if (!isPTOBC) { - if (auto detectedArch = detectTextualModuleArch(buf)) - arch = *detectedArch; - } - if (arch != "a3" && arch != "a5") - arch = "a3"; - - if (isPTOBC) { - // Decode PTO bytecode directly into an MLIR module. - llvm::ArrayRef bytes(reinterpret_cast(buf.data()), buf.size()); -#if defined(__cpp_exceptions) || defined(__EXCEPTIONS) - try { - module = ptobc::decodePTOBCToModule(bytes, context); - } catch (...) { - llvm::errs() << "Error: Failed to decode PTOBC.\n"; - return 1; - } -#else - module = ptobc::decodePTOBCToModule(bytes, context); -#endif - if (!module) { - llvm::errs() << "Error: Failed to decode PTOBC.\n"; - return 1; - } - } else { - // Parse textual MLIR (.pto). - llvm::SourceMgr sourceMgr; - sourceMgr.AddNewSourceBuffer(std::move(*fileOrErr), llvm::SMLoc()); - pto::ScopedPTOParserTargetArch scopedParserArch( - &context, arch == "a5" ? pto::PTOParserTargetArch::A5 - : pto::PTOParserTargetArch::A3); - module = parseSourceFile(sourceMgr, &context); - if (!module) { - llvm::errs() << "Error: Failed to parse MLIR.\n"; - return 1; - } - } - - // If the CLI explicitly requested an arch, it overrides the input module. - // Otherwise, preserve the textual module's arch when present and only fall - // back to the effective default. - if (cliArchSpecified || !module->getOperation()->hasAttr("pto.target_arch")) { - module->getOperation()->setAttr("pto.target_arch", - mlir::StringAttr::get(&context, arch)); - } - - CompilationPlan compilationPlan; - if (failed(buildCompilationPlan(module.get(), cliBackendSpecified, - effectiveBackend, compilationPlan))) - return 1; - effectiveBackend = compilationPlan.effectiveBackend; - - if (effectiveBackend != PTOBackend::VPTO && - compilationPlan.mode != CompilationMode::MixedFatobj && - (emitVPTO || ptoPrintSeamIR || !ptoSeamIRFile.empty())) { - llvm::errs() << "Error: VPTO-specific flags require " - "--pto-backend=vpto or pto.backend = \"vpto\".\n"; - return 1; - } - - PTOBuildLevel effectiveLevel = defaultBuildLevel(); - if (!parseBuildLevel(ptoBuildLevel, effectiveLevel)) { - llvm::errs() << "Error: invalid --pto-level='" << ptoBuildLevel - << "'. Expected 'level1', 'level2', or 'level3'.\n"; - return 1; - } - - bool invalidAutoSyncTailHint = false; - module->walk([&](mlir::func::FuncOp func) { - auto hintAttr = - func->getAttrOfType("pto.auto_sync_tail_hint"); - if (!hintAttr) - return; - - std::string normalizedHint; - if (!parseAutoSyncTailHint(hintAttr.getValue(), normalizedHint)) { - func.emitError("invalid pto.auto_sync_tail_hint '") - << hintAttr.getValue() - << "'. Expected 'barrier-all' (or 'default') or " - "'mte3-to-s-event0'."; - invalidAutoSyncTailHint = true; - return; - } - func->setAttr("pto.auto_sync_tail_hint", - mlir::StringAttr::get(&context, normalizedHint)); - }); - if (invalidAutoSyncTailHint) - return 1; - - bool hasTAssign = false; - module->walk([&](pto::TAssignOp) { hasTAssign = true; }); - - if (hasTAssign && effectiveLevel != PTOBuildLevel::Level3) { - llvm::errs() << "Error: pto.tassign is only supported when " - "--pto-level=level3.\n"; - return 1; - } - - if (hasTAssign && enableInsertSync) { - llvm::errs() << "Error: pto.tassign requires --enable-insert-sync to be " - "disabled.\n"; - return 1; - } - - int enabledAutoSyncModes = - (enableInsertSync ? 1 : 0) + (enableInjectBarrierAllSync ? 1 : 0) + - (enableGraphSyncSolver ? 1 : 0); - if (enabledAutoSyncModes > 1) { - llvm::errs() << "Error: --enable-insert-sync, " - "--enable-inject-barrier-all-sync, and " - "--enable-graph-sync-solver are mutually exclusive.\n"; - return 1; - } - if (hasTAssign && enableInjectBarrierAllSync) { - llvm::errs() << "Error: pto.tassign requires " - "--enable-inject-barrier-all-sync to be disabled.\n"; - return 1; - } - if (hasTAssign && enableGraphSyncSolver) { - llvm::errs() << "Error: pto.tassign requires --enable-graph-sync-solver " - "to be disabled.\n"; - return 1; - } - - if (effectiveLevel == PTOBuildLevel::Level3) { - bool missing = false; - module->walk([&](pto::AllocTileOp op) { - if (!op.getAddr()) { - op.emitError("requires 'addr' operand when --pto-level=level3"); - missing = true; - } - }); - if (missing) - return 1; - } else { - bool hasAddr = false; - module->walk([&](pto::AllocTileOp op) { - if (op.getAddr()) { - op.emitError( - "unexpected 'addr' operand: only supported when --pto-level=level3"); - hasAddr = true; - } - }); - if (hasAddr) - return 1; - } - - if (compilationPlan.mode == CompilationMode::MixedFatobj) - return runMixedFatobjPlan(compilationPlan, effectiveLevel, arch, argc, argv, - outputFilename, llvm::errs()); - - // [Fix] ToolOutputFile Usage - std::error_code ec; - llvm::ToolOutputFile outputFile(outputFilename, ec, llvm::sys::fs::OF_None); - if (ec) { - llvm::errs() << ec.message() << "\n"; - return 1; - } - - const bool hasTileOpsToExpand = hasUnexpandedTileOps(*module); - const bool hasTilelangHelpers = hasTilelangInlineHelpers(*module); - - if (compilationPlan.mode == CompilationMode::VPTO && !hasTileOpsToExpand) { - if (ptoPrintSeamIR || !ptoSeamIRFile.empty()) { - llvm::errs() << "Error: shared pre-backend seam IR is unavailable when " - "skipping the shared PTO-to-VPTO lowering pipeline.\n"; - return 1; - } - if (failed(runVPTOBackendPipeline(module.get(), argc, argv, hasTileOpsToExpand, - hasTilelangHelpers))) - return 1; - return emitVPTOBackendResult(module.get(), outputFile); - } - - // Main PassManager - PassManager pm(&context); - - if (failed(configureSharedPTOPipeline(pm, effectiveLevel))) - return 1; - - llvm::raw_ostream *outputOS = &outputFile.os(); - - if (emitMlirIR) { - if (failed(pm.run(*module))) { - llvm::errs() << "Error: Pass execution failed.\n"; - return 1; - } - module->print(*outputOS); - outputFile.keep(); - return 0; - } - - addSharedPTOSeamPasses(pm); - if (failed(applyConfiguredPassManagerCLOptions(pm, "main PTOAS pipeline"))) - return 1; - - module->getOperation()->setAttr("pto.target_arch", - mlir::StringAttr::get(&context, arch)); - - if (compilationPlan.mode == CompilationMode::VPTO) { - if (failed(pm.run(*module))) { - llvm::errs() << "Error: Pass execution failed.\n"; - return 1; - } - - if (ptoPrintSeamIR) { - module->print(llvm::errs()); - llvm::errs() << "\n"; - } - if (failed(emitSharedPreBackendSeamIR(*module, ptoSeamIRFile))) - return 1; - - if (failed(runVPTOBackendPipeline(module.get(), argc, argv, hasTileOpsToExpand, - hasTilelangHelpers))) - return 1; - return emitVPTOBackendResult(module.get(), outputFile); - } - - if (arch == "a3") { - pm.addPass(pto::createEmitPTOManualPass(pto::PTOArch::A3)); - } else { - pm.addPass(pto::createEmitPTOManualPass(pto::PTOArch::A5)); - } - pm.addPass(emitc::createFormExpressionsPass()); - pm.addPass(mlir::createCSEPass()); - - if (failed(pm.run(*module))) { - llvm::errs() << "Error: Pass execution failed.\n"; - return 1; - } - - std::string cppOutput; - if (failed(emitEmitCSourceArtifact(module.get(), emitAddPtrTrace, cppOutput, - llvm::errs()))) - return 1; - - *outputOS << cppOutput; - outputOS->flush(); - - outputFile.keep(); // Success, keep the file - - return 0; -} diff --git a/tools/ptoas/PTOASDriver.h b/tools/ptoas/PTOASDriver.h deleted file mode 100644 index 61dfbae99a..0000000000 --- a/tools/ptoas/PTOASDriver.h +++ /dev/null @@ -1,18 +0,0 @@ -// 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. - -#ifndef PTOAS_DRIVER_H -#define PTOAS_DRIVER_H - -namespace mlir::pto { - -int runPTOASDriver(int argc, char **argv); - -} // namespace mlir::pto - -#endif diff --git a/tools/ptoas/VPTOHostStubEmission.cpp b/tools/ptoas/VPTOHostStubEmission.cpp index a2698ae3b3..5a557f861b 100644 --- a/tools/ptoas/VPTOHostStubEmission.cpp +++ b/tools/ptoas/VPTOHostStubEmission.cpp @@ -20,10 +20,6 @@ using namespace mlir; namespace { -static bool hasVPTOKernelAttr(Operation *op) { - return op->hasAttr("pto.kernel") || op->hasAttr("pto.aicore"); -} - struct VPTOKernelStubDecl { std::string logicalName; SmallVector argTypes; @@ -81,7 +77,7 @@ static LogicalResult collectVPTOKernelStubDecls( for (ModuleOp module : modules) { module.walk([&](func::FuncOp func) { - if (func.isExternal() || !hasVPTOKernelAttr(func)) + if (func.isExternal() || !pto::isPTOKernelFunction(func)) return; std::string logicalName = getLogicalKernelName(func.getSymName()); diff --git a/tools/ptoas/driver.cpp b/tools/ptoas/driver.cpp new file mode 100644 index 0000000000..cbe93b2eca --- /dev/null +++ b/tools/ptoas/driver.cpp @@ -0,0 +1,822 @@ +// 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. + +#include "ptoas.h" + +#include "ObjectEmission.h" +#include "PTO/IR/PTO.h" +#include "PTO/Transforms/BufferizableOpInterfaceImpl.h" +#include "PTO/Transforms/Passes.h" +#include "VPTOHostStubEmission.h" +#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h" +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Math/IR/Math.h" +#include "mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/EmitC/IR/EmitC.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/SymbolTable.h" +#include "mlir/InitAllPasses.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Parser/Parser.h" +#include "llvm/ADT/StringMap.h" +#include "llvm/Support/CommandLine.h" +#include "llvm/Support/FileSystem.h" +#include "llvm/Support/MemoryBuffer.h" +#include "llvm/Support/Path.h" +#include "llvm/Support/Process.h" +#include "llvm/ADT/ScopeExit.h" +#include "llvm/Support/ToolOutputFile.h" +#include "llvm/Support/SourceMgr.h" +#include "ptobc/ptobc_decode.h" +#include +#include +#include +#include +#include + +using namespace mlir; + +#ifndef PTOAS_RELEASE_VERSION +#define PTOAS_RELEASE_VERSION "unknown" +#endif + +static llvm::cl::opt inputFilename(llvm::cl::Positional, + llvm::cl::desc(""), + llvm::cl::init("-")); + +static llvm::cl::opt outputFilename( + "o", llvm::cl::desc("Output filename"), llvm::cl::value_desc("filename"), + llvm::cl::init("-")); + +static void printPTOASVersion(llvm::raw_ostream &os) { + os << "ptoas " << PTOAS_RELEASE_VERSION << "\n"; +} + +static void registerPTOASDriverDialects(DialectRegistry ®istry) { + registry.insert(); + registry.insert(); + registry.insert(); + registry.insert(); + registry.insert(); + registry.insert(); + registry.insert(); + registry.insert(); + registry.insert(); + + registry.insert(); + arith::registerBufferizableOpInterfaceExternalModels(registry); + tensor::registerBufferizableOpInterfaceExternalModels(registry); + mlir::pto::registerBufferizableOpInterfaceExternalModels(registry); + + registry.insert(); + registry.insert(); +} + +static void registerPTOASDriverPassesAndCLOptions() { + mlir::registerAllPasses(); + mlir::pto::registerPTOPasses(); + mlir::pto::registerPTOViewToMemrefPass(); + mlir::pto::registerPTOInlineLibCall(); + mlir::pto::registerFoldTileBufIntrinsics(); + mlir::pto::registerExpandTileOp(); + mlir::registerPassManagerCLOptions(); +} + +static bool hasCLIOption(int argc, char **argv, llvm::StringRef option) { + const std::string optionWithValue = (option + "=").str(); + for (int i = 1; i < argc; ++i) { + llvm::StringRef arg(argv[i]); + if (arg == option || arg.starts_with(optionWithValue)) + return true; + } + return false; +} + +static void loadPTOASDialects(MLIRContext &context) { + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.getOrLoadDialect(); +} + +static LogicalResult createTempPath(llvm::StringRef prefix, + llvm::StringRef suffix, + std::string &path) { + llvm::SmallString<128> tempPath; + int fd = -1; + std::error_code ec = + llvm::sys::fs::createTemporaryFile(prefix, suffix, fd, tempPath); + if (ec) { + llvm::errs() << "Error: failed to create temporary file for " << prefix + << suffix << ": " << ec.message() << "\n"; + return failure(); + } + llvm::sys::Process::SafelyCloseFileDescriptor(fd); + path = tempPath.str().str(); + return success(); +} + +static void removeTempPaths(ArrayRef paths) { + for (const std::string &path : paths) + if (!path.empty()) + llvm::sys::fs::remove(path); +} + +static std::string sanitizeModuleId(llvm::StringRef outputPath) { + std::string moduleId = + outputPath.empty() || outputPath == "-" ? "ptoas_fatobj" : outputPath.str(); + for (char &c : moduleId) + if (!std::isalnum(static_cast(c)) && c != '_') + c = '_'; + return moduleId; +} + +struct BackendExportPlan { + std::string sourceName; + std::string abiName; + FunctionType type; +}; + +struct BackendImportPlan { + std::string sourceName; + std::string abiName; + FunctionType type; +}; + +struct BackendChildPlan { + ModuleOp module; + OwningOpRef ownedModule; + mlir::pto::PTOASBackend backend = mlir::pto::PTOASBackend::EmitC; + std::optional kernelKind; + SmallVector exports; + SmallVector imports; +}; + +static bool parseDriverBackend(llvm::StringRef backendStr, + mlir::pto::PTOASBackend &out) { + std::string s = backendStr.str(); + for (char &c : s) + c = static_cast(std::tolower(static_cast(c))); + if (s == "emitc") { + out = mlir::pto::PTOASBackend::EmitC; + return true; + } + if (s == "vpto") { + out = mlir::pto::PTOASBackend::VPTO; + return true; + } + return false; +} + +static LogicalResult parseDriverBackendAttr( + Operation *op, std::optional &backend) { + backend = std::nullopt; + Attribute rawBackendAttr = op->getAttr("pto.backend"); + if (!rawBackendAttr) + return success(); + + auto backendAttr = dyn_cast(rawBackendAttr); + if (!backendAttr) { + return op->emitError("invalid pto.backend attribute. Expected string " + "value 'emitc' or 'vpto'."); + } + + mlir::pto::PTOASBackend attrBackend = mlir::pto::PTOASBackend::EmitC; + if (!parseDriverBackend(backendAttr.getValue(), attrBackend)) { + return op->emitError("invalid pto.backend '") + << backendAttr.getValue() << "'. Expected 'emitc' or 'vpto'."; + } + + backend = attrBackend; + return success(); +} + +static LogicalResult resolveDriverModuleBackend( + ModuleOp module, bool cliBackendSpecified, + mlir::pto::PTOASBackend &effectiveBackend, bool &mixedBackendMode) { + mixedBackendMode = false; + if (cliBackendSpecified) + return success(); + + std::optional topBackend; + if (failed(parseDriverBackendAttr(module.getOperation(), topBackend))) + return failure(); + if (topBackend) { + effectiveBackend = *topBackend; + return success(); + } + + std::optional childBackend; + bool sawEmitCChild = false; + bool sawVPTOChild = false; + SmallVector missingBackendChildren; + for (ModuleOp child : module.getOps()) { + std::optional parsedChildBackend; + if (failed(parseDriverBackendAttr(child.getOperation(), parsedChildBackend))) + return failure(); + if (!parsedChildBackend) { + missingBackendChildren.push_back(child); + continue; + } + sawEmitCChild |= *parsedChildBackend == mlir::pto::PTOASBackend::EmitC; + sawVPTOChild |= *parsedChildBackend == mlir::pto::PTOASBackend::VPTO; + + if (!childBackend) { + childBackend = *parsedChildBackend; + continue; + } + if (*childBackend != *parsedChildBackend) { + if (!missingBackendChildren.empty()) { + return missingBackendChildren.front().emitError() + << "mixed-backend container child module is missing pto.backend"; + } + mixedBackendMode = true; + sawEmitCChild = true; + sawVPTOChild = true; + continue; + } + } + + if (sawEmitCChild && sawVPTOChild) + mixedBackendMode = true; + if (childBackend) + effectiveBackend = *childBackend; + if (sawEmitCChild && sawVPTOChild && !missingBackendChildren.empty()) { + return missingBackendChildren.front().emitError() + << "mixed-backend container child module is missing pto.backend"; + } + return success(); +} + +static std::optional +getModuleKernelKind(ModuleOp module) { + auto kernelKind = + module->getAttrOfType( + mlir::pto::FunctionKernelKindAttr::name); + if (!kernelKind) + return std::nullopt; + return kernelKind.getKernelKind(); +} + +static llvm::StringRef +getPublicABISuffix(mlir::pto::FunctionKernelKind kind) { + switch (kind) { + case mlir::pto::FunctionKernelKind::Vector: + return ".vector"; + case mlir::pto::FunctionKernelKind::Cube: + return ".cube"; + } + llvm_unreachable("unknown function kernel kind"); +} + +static OwningOpRef cloneBackendChildModule(ModuleOp outer, + ModuleOp child) { + Operation *clonedOp = child.getOperation()->clone(); + auto clonedModule = cast(clonedOp); + + for (NamedAttribute attr : outer->getAttrs()) { + StringRef attrName = attr.getName().getValue(); + if (attrName == SymbolTable::getSymbolAttrName() || + attrName == "pto.backend") + continue; + if (!clonedModule->hasAttr(attr.getName())) + clonedModule->setAttr(attr.getName(), attr.getValue()); + } + + return OwningOpRef(clonedModule); +} + +static void collectBackendChildSymbols(BackendChildPlan &plan) { + for (func::FuncOp func : plan.module.getOps()) { + llvm::StringRef symName = func.getSymName(); + if (mlir::pto::isPTOKernelFunction(func)) + continue; + + if (func.isDeclaration()) { + BackendImportPlan import; + import.sourceName = symName.str(); + import.type = func.getFunctionType(); + plan.imports.push_back(std::move(import)); + continue; + } + + if (!func.isPublic()) + continue; + + BackendExportPlan exportPlan; + exportPlan.sourceName = symName.str(); + exportPlan.type = func.getFunctionType(); + if (plan.backend == mlir::pto::PTOASBackend::VPTO && plan.kernelKind) + exportPlan.abiName = + (symName + getPublicABISuffix(*plan.kernelKind)).str(); + else + exportPlan.abiName = symName.str(); + plan.exports.push_back(std::move(exportPlan)); + } +} + +static LogicalResult +collectBackendChildren(ModuleOp module, + SmallVectorImpl &children) { + for (ModuleOp child : module.getOps()) { + std::optional childBackend; + if (failed(parseDriverBackendAttr(child.getOperation(), childBackend))) + return failure(); + if (!childBackend) + continue; + + BackendChildPlan plan; + plan.ownedModule = cloneBackendChildModule(module, child); + plan.module = plan.ownedModule.get(); + plan.backend = *childBackend; + plan.kernelKind = getModuleKernelKind(plan.module); + collectBackendChildSymbols(plan); + children.push_back(std::move(plan)); + } + return success(); +} + +static LogicalResult verifyUniqueBackendABIExports( + SmallVectorImpl &children) { + llvm::StringMap seen; + for (BackendChildPlan &child : children) { + for (const BackendExportPlan &exportPlan : child.exports) { + auto inserted = + seen.try_emplace(exportPlan.abiName, child.module.getLoc()); + if (!inserted.second) { + return child.module.emitError() + << "generated ABI export symbol '" << exportPlan.abiName + << "' is defined by multiple backend child modules"; + } + } + } + return success(); +} + +static LogicalResult verifyExternalImportShape( + SmallVectorImpl &children) { + for (BackendChildPlan &child : children) { + for (func::FuncOp func : child.module.getOps()) { + if (!func.isDeclaration()) + continue; + if (!func.isPrivate()) { + return func.emitError() + << "external func.func declaration must use private visibility"; + } + if (mlir::pto::isPTOKernelFunction(func)) { + return func.emitError() + << "external func.func declaration must not be marked pto.aicore"; + } + } + } + return success(); +} + +static LogicalResult resolveExternalImports( + SmallVectorImpl &children) { + for (BackendChildPlan &importingChild : children) { + for (BackendImportPlan &importPlan : importingChild.imports) { + const BackendExportPlan *matchedExport = nullptr; + for (const BackendChildPlan &exportingChild : children) { + if (&exportingChild == &importingChild) + continue; + for (const BackendExportPlan &exportPlan : exportingChild.exports) { + if (exportPlan.sourceName != importPlan.sourceName) + continue; + if (matchedExport) { + return importingChild.module.emitError() + << "cross-backend external import '" << importPlan.sourceName + << "' matches multiple exported definitions"; + } + matchedExport = &exportPlan; + } + } + + if (!matchedExport) { + return importingChild.module.emitError() + << "cross-backend external import '" << importPlan.sourceName + << "' has no matching public non-pto.aicore definition in " + "another backend child module"; + } + + if (matchedExport->type != importPlan.type) { + return importingChild.module.emitError() + << "cross-backend external import '" << importPlan.sourceName + << "' signature does not match exported definition"; + } + + importPlan.abiName = matchedExport->abiName; + } + } + return success(); +} + +static LogicalResult applyResolvedExternalImports(BackendChildPlan &child) { + llvm::StringMap importABIMap; + for (const BackendImportPlan &importPlan : child.imports) { + if (importPlan.abiName.empty()) + return child.module.emitError() + << "unresolved cross-backend external import '" + << importPlan.sourceName << "'"; + importABIMap[importPlan.sourceName] = importPlan.abiName; + } + if (importABIMap.empty()) + return success(); + + for (func::FuncOp func : child.module.getOps()) { + auto it = importABIMap.find(func.getSymName()); + if (it == importABIMap.end()) + continue; + func->setAttr("pto.external_abi", + StringAttr::get(func.getContext(), it->second)); + } + return success(); +} + +static LogicalResult compileMixedBackendContainer( + ModuleOp module, llvm::StringRef arch, bool cliBackendSpecified, int argc, + char **argv, mlir::pto::PTOASCompileResult &result) { + mlir::pto::PTOASBackend effectiveBackend = mlir::pto::PTOASBackend::EmitC; + if (failed(mlir::pto::getPTOASCommandLineBackend(effectiveBackend))) + return failure(); + + bool mixedBackendMode = false; + if (failed(resolveDriverModuleBackend(module, cliBackendSpecified, + effectiveBackend, mixedBackendMode))) + return failure(); + if (!mixedBackendMode) + return mlir::pto::compilePTOASModule(module, arch, cliBackendSpecified, + argc, argv, result); + if (mlir::pto::isPTOASDebugIROutputRequested()) { + llvm::errs() << "Error: mixed pto.backend fatobj mode does not support " + "debug IR output flags.\n"; + return failure(); + } + + SmallVector children; + if (failed(collectBackendChildren(module, children))) + return failure(); + if (failed(verifyUniqueBackendABIExports(children))) + return failure(); + if (failed(verifyExternalImportShape(children))) + return failure(); + if (failed(resolveExternalImports(children))) + return failure(); + + result = mlir::pto::PTOASCompileResult{}; + result.kind = mlir::pto::PTOASCompileResultKind::MixedObject; + SmallVector stubModules; + + for (BackendChildPlan &child : children) { + if (failed(applyResolvedExternalImports(child))) + return failure(); + child.module->setAttr("pto.backend", + StringAttr::get(child.module.getContext(), + child.backend == + mlir::pto::PTOASBackend::VPTO + ? "vpto" + : "emitc")); + + mlir::pto::PTOASCompileResult childResult; + if (failed(mlir::pto::compilePTOASModule(child.module, arch, + /*cliBackendSpecified=*/false, + argc, argv, childResult, + /*emitVPTOHostStub=*/false))) + return failure(); + if (childResult.kind == mlir::pto::PTOASCompileResultKind::Text) { + mlir::pto::PTOASBackendObjectInput object; + object.kind = mlir::pto::PTOASBackendObjectKind::EmitCFatobj; + object.cppSource = std::move(childResult.textOutput); + result.backendObjects.push_back(std::move(object)); + continue; + } + + if (childResult.kind != mlir::pto::PTOASCompileResultKind::VPTOObject) { + llvm::errs() << "Error: nested mixed backend compilation is not " + "supported.\n"; + return failure(); + } + for (mlir::pto::PTOASBackendObjectInput &object : + childResult.backendObjects) + result.backendObjects.push_back(std::move(object)); + if (child.module) + stubModules.push_back(child.module); + } + + if (result.backendObjects.empty()) { + llvm::errs() << "Error: mixed fatobj compilation produced no backend " + "objects.\n"; + return failure(); + } + + if (stubModules.empty()) { + result.stubSource = "#ifndef __global__\n#define __global__\n#endif\n\n" + "#ifndef __gm__\n#define __gm__\n#endif\n\n"; + return success(); + } + if (failed(mlir::pto::emitVPTOHostStubSource(stubModules, result.stubSource, + llvm::errs()))) { + llvm::errs() << "Error: Failed to emit mixed VPTO host stub source.\n"; + return failure(); + } + return success(); +} + +static LogicalResult emitBackendObject( + mlir::pto::PTOASBackendObjectInput &input, + const mlir::pto::ObjectEmissionToolchain &toolchain, + SmallVectorImpl &tempPaths, std::string &outputPath) { + std::string stderrPath; + if (failed(createTempPath("ptoas-object", ".log", stderrPath))) + return failure(); + tempPaths.push_back(stderrPath); + + if (input.kind == mlir::pto::PTOASBackendObjectKind::EmitCFatobj) { + std::string cppPath; + if (failed(createTempPath("ptoas-emitc", ".cpp", cppPath)) || + failed(createTempPath("ptoas-emitc-fatobj", ".o", outputPath))) + return failure(); + tempPaths.push_back(cppPath); + tempPaths.push_back(outputPath); + return mlir::pto::emitCppFatobj(input.cppSource, cppPath, outputPath, + toolchain, stderrPath, llvm::errs()); + } + + std::string llPath; + if (failed(createTempPath("ptoas-vpto", ".ll", llPath)) || + failed(createTempPath("ptoas-vpto-device", ".o", outputPath))) + return failure(); + tempPaths.push_back(llPath); + tempPaths.push_back(outputPath); + + if (!input.llvmModule.module) { + llvm::errs() << "Error: missing LLVM module for VPTO object emission.\n"; + return failure(); + } + + if (input.kind == mlir::pto::PTOASBackendObjectKind::VPTOVectorDeviceObject) { + return mlir::pto::emitVPTOVectorDeviceObject( + *input.llvmModule.module, llPath, outputPath, toolchain, stderrPath, + llvm::errs()); + } + return mlir::pto::emitVPTOCubeDeviceObject( + *input.llvmModule.module, llPath, outputPath, toolchain, stderrPath, + llvm::errs()); +} + +static LogicalResult writeTextOutput(llvm::StringRef output, + llvm::StringRef outputPath) { + std::error_code ec; + llvm::ToolOutputFile outputFile(outputPath, ec, llvm::sys::fs::OF_None); + if (ec) { + llvm::errs() << ec.message() << "\n"; + return failure(); + } + outputFile.os() << output; + outputFile.os().flush(); + outputFile.keep(); + return success(); +} + +static LogicalResult emitObjectOutput(mlir::pto::PTOASCompileResult &result, + llvm::StringRef outputPath) { + if (outputPath.empty() || outputPath == "-") { + if (result.kind == mlir::pto::PTOASCompileResultKind::MixedObject) { + llvm::errs() << "Error: mixed pto.backend fatobj mode requires an " + "explicit file path passed with -o.\n"; + return failure(); + } + llvm::errs() << "Error: object output requires an explicit file path " + "passed with -o.\n"; + return failure(); + } + + mlir::pto::ObjectEmissionToolchain toolchain; + if (failed(mlir::pto::discoverObjectEmissionToolchain(toolchain, + llvm::errs()))) + return failure(); + + SmallVector tempPaths; + auto cleanup = llvm::make_scope_exit([&]() { removeTempPaths(tempPaths); }); + + std::string stubPath; + std::string stderrPath; + if (failed(createTempPath("ptoas-stub", ".cpp", stubPath)) || + failed(createTempPath("ptoas-fatobj", ".log", stderrPath))) + return failure(); + tempPaths.push_back(stubPath); + tempPaths.push_back(stderrPath); + if (failed(mlir::pto::writeHostStubSource(result.stubSource, stubPath, + llvm::errs()))) + return failure(); + + SmallVector fatobjPaths; + const std::string moduleId = sanitizeModuleId(outputPath); + if (result.kind == mlir::pto::PTOASCompileResultKind::VPTOObject) { + SmallVector deviceObjectPaths; + for (size_t i = 0, e = result.backendObjects.size(); i < e; ++i) { + std::string objectPath; + if (failed(emitBackendObject(result.backendObjects[i], toolchain, + tempPaths, objectPath))) + return failure(); + deviceObjectPaths.push_back(objectPath); + } + + std::string mergedObjPath; + std::string fatobjPath; + if (failed(createTempPath("ptoas-device-merged", ".o", mergedObjPath)) || + failed(createTempPath("ptoas-device-fatobj", ".o", fatobjPath))) + return failure(); + tempPaths.push_back(mergedObjPath); + tempPaths.push_back(fatobjPath); + + if (failed(mlir::pto::mergeDeviceObjects(deviceObjectPaths, mergedObjPath, + toolchain, stderrPath, + llvm::errs()))) + return failure(); + if (failed(mlir::pto::compileStubToFatobj( + stubPath, mergedObjPath, fatobjPath, moduleId, toolchain, + stderrPath, llvm::errs()))) + return failure(); + + if (std::error_code ec = llvm::sys::fs::copy_file(fatobjPath, outputPath)) { + llvm::errs() << "Error: failed to copy fatobj to " << outputPath << ": " + << ec.message() << "\n"; + return failure(); + } + return success(); + } + + for (size_t i = 0, e = result.backendObjects.size(); i < e; ++i) { + std::string objectPath; + if (failed(emitBackendObject(result.backendObjects[i], toolchain, tempPaths, + objectPath))) + return failure(); + + if (result.backendObjects[i].kind == + mlir::pto::PTOASBackendObjectKind::EmitCFatobj) { + fatobjPaths.push_back(objectPath); + continue; + } + + std::string mergedObjPath; + std::string fatobjPath; + if (failed(createTempPath("ptoas-device-merged", ".o", mergedObjPath)) || + failed(createTempPath("ptoas-device-fatobj", ".o", fatobjPath))) + return failure(); + tempPaths.push_back(mergedObjPath); + tempPaths.push_back(fatobjPath); + + if (failed(mlir::pto::mergeDeviceObjects(ArrayRef(objectPath), + mergedObjPath, toolchain, + stderrPath, llvm::errs()))) + return failure(); + + if (failed(mlir::pto::compileStubToFatobj( + stubPath, mergedObjPath, fatobjPath, + moduleId + "_backend_" + std::to_string(i), toolchain, stderrPath, + llvm::errs()))) + return failure(); + fatobjPaths.push_back(fatobjPath); + } + + if (fatobjPaths.empty()) { + llvm::errs() << "Error: object emission produced no fatobjs.\n"; + return failure(); + } + + if (fatobjPaths.size() == 1) { + if (std::error_code ec = llvm::sys::fs::copy_file(fatobjPaths.front(), + outputPath)) { + llvm::errs() << "Error: failed to copy fatobj to " << outputPath << ": " + << ec.message() << "\n"; + return failure(); + } + return success(); + } + + if (failed(mlir::pto::linkFatobjs(fatobjPaths, outputPath, toolchain, + stderrPath, llvm::errs()))) + return failure(); + return success(); +} + +int main(int argc, char **argv) { + DialectRegistry registry; + registerPTOASDriverDialects(registry); + registerPTOASDriverPassesAndCLOptions(); + llvm::cl::SetVersionPrinter(printPTOASVersion); + + const bool cliArchSpecified = hasCLIOption(argc, argv, "--pto-arch"); + const bool cliBackendSpecified = hasCLIOption(argc, argv, "--pto-backend"); + + llvm::cl::ParseCommandLineOptions(argc, argv, "PTO Assembler (ptoas)\n"); + + auto fileOrErr = llvm::MemoryBuffer::getFileOrSTDIN(inputFilename); + if (!fileOrErr) { + llvm::errs() << "Error: Could not open input file: " + << fileOrErr.getError().message() << "\n"; + return 1; + } + + MLIRContext context(registry); + // Be tolerant: ptobc decode may materialize ops from dialects that aren't + // explicitly registered/loaded in this tool yet. + context.allowUnregisteredDialects(true); + loadPTOASDialects(context); + + OwningOpRef module; + llvm::StringRef buf = (*fileOrErr)->getBuffer(); + const bool isPTOBC = + (buf.size() >= 6 && std::memcmp(buf.data(), "PTOBC\0", 6) == 0); + + std::string arch = + mlir::pto::normalizePTOASArch(mlir::pto::getPTOASTargetArchOption()); + if (cliArchSpecified) { + if (!mlir::pto::isSupportedPTOASArch(arch)) { + llvm::errs() << "Error: invalid --pto-arch='" + << mlir::pto::getPTOASTargetArchOption() + << "'. Expected 'a3' or 'a5'.\n"; + return 1; + } + } else if (!isPTOBC) { + if (auto detectedArch = mlir::pto::detectPTOASTextualModuleArch(buf)) + arch = *detectedArch; + } + if (!mlir::pto::isSupportedPTOASArch(arch)) + arch = "a3"; + + if (isPTOBC) { + llvm::ArrayRef bytes( + reinterpret_cast(buf.data()), buf.size()); +#if defined(__cpp_exceptions) || defined(__EXCEPTIONS) + try { + module = ptobc::decodePTOBCToModule(bytes, context); + } catch (...) { + llvm::errs() << "Error: Failed to decode PTOBC.\n"; + return 1; + } +#else + module = ptobc::decodePTOBCToModule(bytes, context); +#endif + if (!module) { + llvm::errs() << "Error: Failed to decode PTOBC.\n"; + return 1; + } + } else { + llvm::SourceMgr sourceMgr; + sourceMgr.AddNewSourceBuffer(std::move(*fileOrErr), llvm::SMLoc()); + mlir::pto::ScopedPTOParserTargetArch scopedParserArch( + &context, arch == "a5" ? mlir::pto::PTOParserTargetArch::A5 + : mlir::pto::PTOParserTargetArch::A3); + module = parseSourceFile(sourceMgr, &context); + if (!module) { + llvm::errs() << "Error: Failed to parse MLIR.\n"; + return 1; + } + } + + if (cliArchSpecified || !module->getOperation()->hasAttr("pto.target_arch")) { + module->getOperation()->setAttr("pto.target_arch", + mlir::StringAttr::get(&context, arch)); + } + + mlir::pto::PTOASBackend effectiveBackend = mlir::pto::PTOASBackend::EmitC; + bool mixedBackendMode = false; + if (failed(mlir::pto::getPTOASCommandLineBackend(effectiveBackend)) || + failed(resolveDriverModuleBackend(module.get(), cliBackendSpecified, + effectiveBackend, mixedBackendMode))) + return 1; + if (mixedBackendMode && (outputFilename.empty() || outputFilename == "-")) { + llvm::errs() << "Error: mixed pto.backend fatobj mode requires an " + "explicit file path passed with -o.\n"; + return 1; + } + + mlir::pto::PTOASCompileResult result; + if (failed(compileMixedBackendContainer(module.get(), arch, + cliBackendSpecified, argc, argv, + result))) + return 1; + + if (result.kind == mlir::pto::PTOASCompileResultKind::Text) + return failed(writeTextOutput(result.textOutput, outputFilename)); + + return failed(emitObjectOutput(result, outputFilename)); +} diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index b4b707c010..4f5cdd7506 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -6,8 +6,1727 @@ // 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. -#include "PTOASDriver.h" +#include "ptoas.h" -int main(int argc, char **argv) { - return mlir::pto::runPTOASDriver(argc, argv); +#include "PTO/IR/PTO.h" +#include "PTO/Transforms/VPTOLLVMEmitter.h" +#include "PTO/Transforms/Passes.h" +#include "PTO/Transforms/BufferizableOpInterfaceImpl.h" +#include "VPTOHostStubEmission.h" +#include "TilelangDaemon.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/Diagnostics.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/InitAllDialects.h" +#include "mlir/InitAllPasses.h" +#include "mlir/Parser/Parser.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Math/IR/Math.h" +#include +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Target/Cpp/CppEmitter.h" +#include "llvm/Support/SourceMgr.h" +#include "llvm/Support/ToolOutputFile.h" +#include "llvm/Support/FileSystem.h" // [Fix] Required for OF_None +#include "llvm/Support/Path.h" +#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/EmitC/IR/EmitC.h" +#include "mlir/Dialect/EmitC/Transforms/Passes.h" +#include "llvm/Support/CommandLine.h" +#include "llvm/Support/Regex.h" +#include "llvm/Support/raw_ostream.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/ADT/StringSwitch.h" +#include "llvm/ADT/StringMap.h" +#include +#include +#include +#include +#include +#include +#include + +using namespace mlir; +using namespace pto; + +#ifndef PTOAS_RELEASE_VERSION +#define PTOAS_RELEASE_VERSION "unknown" +#endif + +static std::string getParentDir(llvm::StringRef path) { + llvm::SmallString<256> parent(path); + llvm::sys::path::remove_filename(parent); + llvm::sys::path::remove_dots(parent, true); + return std::string(parent); +} + +static bool pathExists(llvm::StringRef path) { + return !path.empty() && llvm::sys::fs::exists(path); +} + +static std::string joinPath(llvm::StringRef lhs, llvm::StringRef rhs) { + llvm::SmallString<256> joined(lhs); + llvm::sys::path::append(joined, rhs); + llvm::sys::path::remove_dots(joined, true); + return std::string(joined); +} + +static std::string detectInstalledTilelangPath(const char *argv0) { + std::string exePath = + llvm::sys::fs::getMainExecutable(argv0, (void *)&compilePTOASModule); + if (exePath.empty()) + return {}; + + const std::string exeDir = getParentDir(exePath); + const std::string prefixDir = getParentDir(exeDir); + const std::string installedTileOps = joinPath(prefixDir, "share/ptoas/TileOps"); + if (pathExists(installedTileOps)) + return installedTileOps; + return {}; +} + +static std::string detectInstalledTilelangPkgPath(const char *argv0) { + std::string exePath = + llvm::sys::fs::getMainExecutable(argv0, (void *)&compilePTOASModule); + if (exePath.empty()) + return {}; + + const std::string exeDir = getParentDir(exePath); + const std::string prefixDir = getParentDir(exeDir); + const std::string installedPkgRoot = prefixDir; + const std::string installedPkg = joinPath(installedPkgRoot, "tilelang_dsl"); + if (pathExists(installedPkg)) + return installedPkgRoot; + return {}; +} + +static bool hasCLIOption(int argc, char **argv, llvm::StringRef option) { + const std::string optionWithValue = (option + "=").str(); + for (int i = 1; i < argc; ++i) { + llvm::StringRef arg(argv[i]); + if (arg == option || arg.starts_with(optionWithValue)) + return true; + } + return false; +} + +static LogicalResult applyConfiguredPassManagerCLOptions( + PassManager &pm, llvm::StringRef pipelineName, + llvm::raw_ostream &diagOS = llvm::errs()) { + if (succeeded(mlir::applyPassManagerCLOptions(pm))) + return success(); + diagOS << "Error: failed to apply MLIR pass manager command-line options for " + << pipelineName << ".\n"; + return failure(); +} + +static LogicalResult reorderEmitCFunctions(ModuleOp module) { + SmallVector declarations; + SmallVector definitions; + llvm::DenseMap definitionsByName; + + for (auto func : module.getOps()) { + if (func.isDeclaration()) { + declarations.push_back(func); + continue; + } + definitions.push_back(func); + definitionsByName[func.getSymNameAttr()] = func; + } + + llvm::DenseMap indegree; + llvm::DenseMap> outgoing; + for (auto func : definitions) + indegree[func.getOperation()] = 0; + + for (auto caller : definitions) { + Operation *callerOp = caller.getOperation(); + llvm::SmallPtrSet seenCallees; + bool hasCycle = false; + caller.walk([&](emitc::CallOp call) { + auto calleeAttr = call.getCalleeAttr(); + if (!calleeAttr) + return; + auto it = definitionsByName.find(calleeAttr.getLeafReference()); + if (it == definitionsByName.end()) + return; + Operation *calleeOp = it->second.getOperation(); + if (calleeOp == callerOp) { + hasCycle = true; + return; + } + if (!seenCallees.insert(calleeOp).second) + return; + outgoing[calleeOp].push_back(callerOp); + ++indegree[callerOp]; + }); + if (hasCycle) { + return caller.emitOpError() + << "recursive function calls are not supported for EmitC C++ " + "emission"; + } + } + + SmallVector ready; + for (auto func : definitions) { + if (indegree[func.getOperation()] == 0) + ready.push_back(func.getOperation()); + } + + SmallVector sortedDefinitions; + while (!ready.empty()) { + Operation *next = ready.front(); + ready.erase(ready.begin()); + auto nextFunc = cast(next); + sortedDefinitions.push_back(nextFunc); + + for (Operation *user : outgoing[next]) { + unsigned &userIndegree = indegree[user]; + if (--userIndegree == 0) + ready.push_back(user); + } + } + + if (sortedDefinitions.size() != definitions.size()) { + return module.emitError() + << "cyclic function call graph is not supported for EmitC C++ emission"; + } + + if (declarations.empty() && definitions.size() <= 1) + return success(); + + SmallVector desiredOrder; + desiredOrder.append(declarations.begin(), declarations.end()); + desiredOrder.append(sortedDefinitions.begin(), sortedDefinitions.end()); + + Block &body = module.getBodyRegion().front(); + Operation *anchor = nullptr; + for (Operation &op : body.getOperations()) { + if (isa(op)) { + anchor = &op; + break; + } + } + if (!anchor) + return success(); + + auto advanceAnchor = [&]() { + while (anchor) { + anchor = anchor->getNextNode(); + if (!anchor || isa(anchor)) + return; + } + }; + + for (auto func : desiredOrder) { + if (func.getOperation() == anchor) { + advanceAnchor(); + continue; + } + if (anchor) + func->moveBefore(anchor); + else + func->moveBefore(&body, body.end()); + } + + return success(); +} + +static llvm::cl::opt enableInsertSync("enable-insert-sync", + llvm::cl::desc("Enable automatic synchronization insertion pass"), + llvm::cl::init(false)); + +static llvm::cl::opt enableInjectBarrierAllSync( + "enable-inject-barrier-all-sync", + llvm::cl::desc("Enable conservative synchronization by inserting " + "pto.barrier PIPE_ALL before memory-effecting PTO pipe ops"), + llvm::cl::init(false)); + +static llvm::cl::opt enableGraphSyncSolver( + "enable-graph-sync-solver", + llvm::cl::desc("Enable the graph-based intra-core sync solver " + "(experimental). Mutually exclusive with " + "--enable-insert-sync and " + "--enable-inject-barrier-all-sync."), + llvm::cl::init(false)); + +static llvm::cl::opt graphSyncSolverEventIdMax( + "graph-sync-solver-event-id-max", + llvm::cl::desc( + "Maximum EVENT_ID slots for the graph sync solver (default 8). " + "Lower values exercise the PIPE_ALL coloring fallback sooner."), + llvm::cl::init(8)); + +static llvm::cl::opt enableTileOpExpand( + "enable-tile-op-expand", + llvm::cl::desc( + "Deprecated compatibility flag. TileOp expansion is controlled by " + "--pto-backend=vpto."), + llvm::cl::init(false)); + +#ifndef PTOAS_DEFAULT_TILELANG_PATH +#define PTOAS_DEFAULT_TILELANG_PATH "" +#endif +#ifndef PTOAS_DEFAULT_TILELANG_PKG_PATH +#define PTOAS_DEFAULT_TILELANG_PKG_PATH "" +#endif + +static llvm::cl::opt tilelangPath( + "tilelang-path", + llvm::cl::desc("Path to directory of .py tilelang DSL template files " + "(default: /lib/TileOps, baked in at build time)"), + llvm::cl::init(PTOAS_DEFAULT_TILELANG_PATH)); + +static llvm::cl::opt tilelangPkgPath( + "tilelang-pkg-path", + llvm::cl::desc("PYTHONPATH for tilelang_dsl package " + "(default: /tilelang-dsl/python, baked in at build time)"), + llvm::cl::init(PTOAS_DEFAULT_TILELANG_PKG_PATH)); + +static llvm::cl::opt daemonSocketPath( + "daemon-socket-path", + llvm::cl::desc("Path to Unix domain socket for daemon RPC " + "(default: /tmp/tilelang_daemon_{pid}.sock)"), + llvm::cl::init("")); + +static pto::ExpandTileOpOptions resolveExpandTileOpOptions(int argc, + char **argv) { + pto::ExpandTileOpOptions expandOpts; + expandOpts.tilelangPath = tilelangPath; + expandOpts.tilelangPkgPath = tilelangPkgPath; + + if (!hasCLIOption(argc, argv, "--tilelang-path")) { + std::string detectedTilelangPath = detectInstalledTilelangPath(argv[0]); + if (!detectedTilelangPath.empty()) + expandOpts.tilelangPath = detectedTilelangPath; + } + + if (!hasCLIOption(argc, argv, "--tilelang-pkg-path")) { + std::string detectedTilelangPkgPath = detectInstalledTilelangPkgPath(argv[0]); + if (!detectedTilelangPkgPath.empty()) + expandOpts.tilelangPkgPath = detectedTilelangPkgPath; + } + + // Daemon mode is default (no CLI option needed) + // Automatically start daemon for instance caching + if (!expandOpts.tilelangPath.empty()) { + std::string socket = daemonSocketPath; + if (socket.empty()) + socket = ptoas::DaemonManager::generateSocketPath(); + + // Register cleanup handler (daemon will be stopped on PTOAS exit) + ptoas::registerDaemonCleanup(); + + // Try to start daemon automatically + if (ptoas::DaemonManager::start(socket, expandOpts.tilelangPath, expandOpts.tilelangPkgPath)) { + expandOpts.daemonSocketPath = socket; + llvm::errs() << "Info: TileLang daemon started successfully\n"; + } else { + // Fallback: daemon failed, use subprocess mode (current approach) + expandOpts.daemonSocketPath = ""; + llvm::errs() << "Warning: Failed to start daemon, using subprocess mode (fallback)\n"; + } + } + + return expandOpts; +} + +static llvm::cl::opt disableInferLayout( + "disable-infer-layout", + llvm::cl::desc("Disable PTO layout inference pass (static-only)"), + llvm::cl::init(false)); + +static llvm::cl::opt emitAddPtrTrace( + "emit-addptr-trace", + llvm::cl::desc("Emit addptr trace comments in generated C++ output"), + llvm::cl::init(false)); + +static llvm::cl::opt emitMlirIR( + "emit-pto-ir", + llvm::cl::desc("Emit PTO IR after lowering instead of C++"), + llvm::cl::init(false)); + +static llvm::cl::opt ptoTargetArch( + "pto-arch", + llvm::cl::desc("Target Ascend architecture for codegen: a3 or a5 (default: a3)"), + llvm::cl::value_desc("a3|a5"), + llvm::cl::init("a3")); + +static llvm::cl::opt ptoBuildLevel( + "pto-level", + llvm::cl::desc("Build level for pass pipeline: level1, level2, or level3 (default: level2)"), + llvm::cl::value_desc("level1|level2|level3"), + llvm::cl::init("level2")); + +static llvm::cl::opt ptoBackend( + "pto-backend", + llvm::cl::desc("Final PTOAS backend: emitc or vpto (default: emitc)"), + llvm::cl::value_desc("emitc|vpto"), llvm::cl::init("emitc")); + +static llvm::cl::opt emitVPTO( + "emit-vpto", + llvm::cl::desc("Write final post-pass VPTO IR to -o"), + llvm::cl::init(false)); + +static llvm::cl::opt vptoPrintIR( + "vpto-print-ir", + llvm::cl::desc("Print post-pass VPTO backend IR to stderr"), + llvm::cl::init(false)); + +static llvm::cl::opt vptoLoweringStrategy( + "vpto-lowering-strategy", + llvm::cl::desc("VPTO vector lowering strategy: post-update or no-post-update"), + llvm::cl::value_desc("post-update|no-post-update"), + llvm::cl::init("post-update")); + +static llvm::cl::opt dumpVPTOIR( + "dump-vpto-ir", + llvm::cl::desc("Print post-pass VPTO backend IR to stderr"), + llvm::cl::init(false)); + +static llvm::cl::opt ptoPrintSeamIR( + "pto-print-seam-ir", + llvm::cl::desc("Print shared pre-backend seam IR to stderr"), + llvm::cl::init(false)); + +static llvm::cl::opt ptoSeamIRFile( + "pto-seam-ir-file", + llvm::cl::desc("Write shared pre-backend seam IR to a file"), + llvm::cl::value_desc("path"), + llvm::cl::init("")); + +enum class PTOBuildLevel { + Level1, + Level2, + Level3, +}; + +static PTOBuildLevel defaultBuildLevel() { + return PTOBuildLevel::Level2; +} + +llvm::StringRef mlir::pto::getPTOASTargetArchOption() { + return ptoTargetArch; +} + +bool mlir::pto::isPTOASDebugIROutputRequested() { + return emitMlirIR || emitVPTO || ptoPrintSeamIR || !ptoSeamIRFile.empty(); +} + +std::string mlir::pto::normalizePTOASArch(llvm::StringRef archValue) { + std::string normalized = archValue.str(); + for (char &c : normalized) + c = static_cast(std::tolower(static_cast(c))); + return normalized; +} + +bool mlir::pto::isSupportedPTOASArch(llvm::StringRef archValue) { + return archValue == "a3" || archValue == "a5"; +} + +std::optional +mlir::pto::detectPTOASTextualModuleArch(llvm::StringRef text) { + llvm::SmallVector matches; + llvm::Regex archRegex( + R"ptoarch("?(pto\.target_arch)"?[[:space:]]*=[[:space:]]*"([[:alpha:][:digit:]_]+)")ptoarch"); + if (!archRegex.match(text, &matches) || matches.size() < 3) + return std::nullopt; + return normalizePTOASArch(matches[2]); +} + +static bool parseBuildLevel(llvm::StringRef levelStr, PTOBuildLevel &out) { + std::string s = levelStr.str(); + for (char &c : s) + c = static_cast(std::tolower(static_cast(c))); + if (s == "level1") { + out = PTOBuildLevel::Level1; + return true; + } + if (s == "level2") { + out = PTOBuildLevel::Level2; + return true; + } + if (s == "level3") { + out = PTOBuildLevel::Level3; + return true; + } + return false; +} + +static constexpr llvm::StringLiteral kAutoSyncTailPolicyBarrierAll = + "barrier_all"; +static constexpr llvm::StringLiteral kAutoSyncTailPolicyMte3ToSEvent0 = + "setwait_mte3_to_s_event0"; + +static bool parseAutoSyncTailHint(llvm::StringRef hintStr, std::string &normalized) { + std::string s = hintStr.str(); + for (char &c : s) + c = static_cast(std::tolower(static_cast(c))); + if (s == "barrier-all" || s == "barrier_all" || s == "default") { + normalized = kAutoSyncTailPolicyBarrierAll.str(); + return true; + } + if (s == "mte3-to-s-event0" || s == "mte3_to_s_event0" || + s == "setwait-mte3-to-s-event0" || + s == "setwait_mte3_to_s_event0") { + normalized = kAutoSyncTailPolicyMte3ToSEvent0.str(); + return true; + } + return false; +} + +static bool parseBackend(llvm::StringRef backendStr, PTOASBackend &out) { + std::string s = backendStr.str(); + for (char &c : s) + c = static_cast(std::tolower(static_cast(c))); + if (s == "emitc") { + out = PTOASBackend::EmitC; + return true; + } + if (s == "vpto") { + out = PTOASBackend::VPTO; + return true; + } + return false; +} + +LogicalResult mlir::pto::getPTOASCommandLineBackend(PTOASBackend &backend) { + if (!parseBackend(ptoBackend, backend)) { + llvm::errs() << "Error: invalid --pto-backend='" << ptoBackend + << "'. Expected 'emitc' or 'vpto'.\n"; + return failure(); + } + return success(); +} + +static LogicalResult parsePTOBackendAttr(Operation *op, + std::optional &backend) { + backend = std::nullopt; + Attribute rawBackendAttr = op->getAttr("pto.backend"); + if (!rawBackendAttr) + return success(); + + auto backendAttr = dyn_cast(rawBackendAttr); + if (!backendAttr) { + return op->emitError("invalid pto.backend attribute. Expected string " + "value 'emitc' or 'vpto'."); + } + + PTOASBackend attrBackend = PTOASBackend::EmitC; + if (!parseBackend(backendAttr.getValue(), attrBackend)) { + return op->emitError("invalid pto.backend '") + << backendAttr.getValue() << "'. Expected 'emitc' or 'vpto'."; + } + + backend = attrBackend; + return success(); +} + +static LogicalResult resolveModuleBackend(ModuleOp module, + bool cliBackendSpecified, + bool &mixedBackendMode, + PTOASBackend &effectiveBackend) { + mixedBackendMode = false; + if (cliBackendSpecified) + return success(); + + std::optional topBackend; + if (failed(parsePTOBackendAttr(module.getOperation(), topBackend))) + return failure(); + if (topBackend) { + effectiveBackend = *topBackend; + return success(); + } + + std::optional childBackend; + bool sawEmitCChild = false; + bool sawVPTOChild = false; + SmallVector missingBackendChildren; + for (ModuleOp child : module.getOps()) { + std::optional parsedChildBackend; + if (failed(parsePTOBackendAttr(child.getOperation(), parsedChildBackend))) + return failure(); + if (!parsedChildBackend) { + missingBackendChildren.push_back(child); + continue; + } + sawEmitCChild |= *parsedChildBackend == PTOASBackend::EmitC; + sawVPTOChild |= *parsedChildBackend == PTOASBackend::VPTO; + + if (!childBackend) { + childBackend = *parsedChildBackend; + continue; + } + if (*childBackend != *parsedChildBackend) { + if (!missingBackendChildren.empty()) { + return missingBackendChildren.front().emitError() + << "mixed-backend container child module is missing pto.backend"; + } + mixedBackendMode = true; + sawEmitCChild = true; + sawVPTOChild = true; + continue; + } + } + + if (sawEmitCChild && sawVPTOChild) + mixedBackendMode = true; + if (childBackend) + effectiveBackend = *childBackend; + if (sawEmitCChild && sawVPTOChild && !missingBackendChildren.empty()) { + return missingBackendChildren.front().emitError() + << "mixed-backend container child module is missing pto.backend"; + } + return success(); +} + +enum class CompilationMode { + EmitC, + VPTO, +}; + +static std::optional getModuleKernelKind(ModuleOp module) { + auto kernelKind = module->getAttrOfType( + FunctionKernelKindAttr::name); + if (!kernelKind) + return std::nullopt; + return kernelKind.getKernelKind(); +} + +static llvm::StringRef getPublicABISuffix(FunctionKernelKind kind) { + switch (kind) { + case FunctionKernelKind::Vector: + return ".vector"; + case FunctionKernelKind::Cube: + return ".cube"; + } + llvm_unreachable("unknown function kernel kind"); +} + +static LogicalResult applyVPTOPublicABIExportNames(ModuleOp module) { + LogicalResult status = success(); + module.walk([&](ModuleOp child) { + if (failed(status)) + return WalkResult::interrupt(); + std::optional kernelKind = getModuleKernelKind(child); + if (!kernelKind) + return WalkResult::advance(); + + llvm::StringMap renameMap; + for (func::FuncOp func : child.getOps()) { + if (!func.isPublic() || func.isDeclaration() || + pto::isPTOKernelFunction(func)) + continue; + llvm::StringRef oldName = func.getSymName(); + std::string newName = (oldName + getPublicABISuffix(*kernelKind)).str(); + renameMap[oldName] = newName; + } + + for (func::FuncOp func : child.getOps()) { + func.walk([&](func::CallOp call) { + auto it = renameMap.find(call.getCallee()); + if (it != renameMap.end()) + call.setCallee(it->second); + }); + } + + for (func::FuncOp func : child.getOps()) { + auto it = renameMap.find(func.getSymName()); + if (it != renameMap.end()) + func.setSymName(it->second); + } + return WalkResult::advance(); + }); + return status; +} + +static LogicalResult verifyPTOASBackendInput(ModuleOp module, + bool mixedBackendMode) { + if (mixedBackendMode) { + return module.emitError() + << "mixed pto.backend containers must be handled by the PTOAS driver"; + } + return success(); +} + +static LogicalResult emitSharedPreBackendSeamIR(ModuleOp module, + llvm::StringRef outputPath) { + if (outputPath.empty()) + return success(); + + if (outputPath == "-") { + module->print(llvm::outs()); + llvm::outs() << "\n"; + llvm::outs().flush(); + return success(); + } + + std::error_code ec; + llvm::ToolOutputFile outputFile(outputPath, ec, llvm::sys::fs::OF_None); + if (ec) { + llvm::errs() << "Error: failed to open seam IR file '" << outputPath + << "': " << ec.message() << "\n"; + return failure(); + } + + module->print(outputFile.os()); + outputFile.os() << "\n"; + outputFile.keep(); + return success(); +} + +static bool hasUnexpandedTileOps(ModuleOp module) { + bool found = false; + module.walk([&](Operation *op) { + if (found) + return; + if (isa(op)) + found = true; + }); + return found; +} + +static bool hasTilelangInlineHelpers(ModuleOp module) { + bool found = false; + module.walk([&](func::FuncOp func) { + if (found) + return; + if (func->hasAttr("pto.tilelang.inline_proc")) + found = true; + }); + return found; +} + +// -------------------------------------------------------------------------- +// Post-process C++ output: rewrite marker calls into Tile member calls. +// +// We emit marker calls in EmitC IR because EmitC currently does not provide a +// first-class op for member-function invocation. After translation, we rewrite: +// PTOAS__TILE_SET_VALUE(dst, offset, val) -> dst.SetValue(offset, val) +// PTOAS__TILE_GET_VALUE(src, offset) -> src.GetValue(offset) +// PTOAS__TILE_DATA(obj) -> obj.data() +// PTOAS__TILE_SET_VALIDSHAPE(obj, r, c) -> obj.SetValidShape(r, c) +// PTOAS__PTR_LOAD(ptr, offset) -> ptr[offset] +// PTOAS__PTR_STORE(ptr, offset, val) -> ptr[offset] = val +// PTOAS__EVENTID_ARRAY_LOAD(arr, idx) -> arr[idx] +// PTOAS__EVENTID_ARRAY_STORE(arr, idx, v) -> arr[idx] = v +// -------------------------------------------------------------------------- +struct ParsedMarkerCall { + size_t markerPos = std::string::npos; + size_t rparenPos = std::string::npos; + llvm::SmallVector args; +}; + +struct MarkerRewriteSpec { + llvm::StringRef marker; + llvm::StringRef memberName; + unsigned expectedNumArgs = 0; +}; + +struct MarkerSubscriptRewriteSpec { + llvm::StringRef marker; + unsigned expectedNumArgs = 0; + bool isStore = false; +}; + +static bool parseMarkerArgs(llvm::StringRef argsRef, + llvm::SmallVectorImpl &args) { + size_t partBegin = 0; + int parenDepth = 0; + for (size_t i = 0; i < argsRef.size(); ++i) { + char c = argsRef[i]; + if (c == '(') { + ++parenDepth; + continue; + } + if (c == ')') { + if (parenDepth > 0) + --parenDepth; + continue; + } + if (c == ',' && parenDepth == 0) { + args.push_back(argsRef.slice(partBegin, i).trim()); + partBegin = i + 1; + } + } + if (partBegin > argsRef.size()) + return false; + args.push_back(argsRef.drop_front(partBegin).trim()); + return true; +} + +static std::optional +findNextMarkerCall(const std::string &cpp, llvm::StringRef marker, + size_t searchPos) { + ParsedMarkerCall call; + call.markerPos = cpp.find(marker.str(), searchPos); + if (call.markerPos == std::string::npos) + return std::nullopt; + + size_t lparenPos = call.markerPos + marker.size(); + if (lparenPos >= cpp.size() || cpp[lparenPos] != '(') + return ParsedMarkerCall{call.markerPos, std::string::npos, {}}; + + size_t argsBegin = lparenPos + 1; + int parenDepth = 0; + for (size_t i = argsBegin; i < cpp.size(); ++i) { + char c = cpp[i]; + if (c == '(') { + ++parenDepth; + continue; + } + if (c != ')') + continue; + if (parenDepth == 0) { + call.rparenPos = i; + break; + } + --parenDepth; + } + if (call.rparenPos == std::string::npos) + return call; + + llvm::StringRef argsRef(cpp.data() + argsBegin, call.rparenPos - argsBegin); + if (!parseMarkerArgs(argsRef, call.args)) + call.args.clear(); + return call; +} + +template +static bool rewriteMarkerCalls(std::string &cpp, llvm::StringRef marker, + BuildReplacementFn buildReplacement) { + size_t searchPos = 0; + bool changed = false; + for (auto call = findNextMarkerCall(cpp, marker, searchPos); call; + call = findNextMarkerCall(cpp, marker, searchPos)) { + if (call->rparenPos == std::string::npos) { + searchPos = call->markerPos + marker.size(); + continue; + } + + std::optional replacement = buildReplacement(*call); + if (!replacement) { + searchPos = call->rparenPos + 1; + continue; + } + + cpp.replace(call->markerPos, (call->rparenPos - call->markerPos) + 1, + *replacement); + changed = true; + searchPos = call->markerPos + replacement->size(); + } + return changed; +} + +static bool rewriteMarkerCallToMember(std::string &cpp, llvm::StringRef marker, + llvm::StringRef memberName, + unsigned expectedNumArgs) { + return rewriteMarkerCalls( + cpp, marker, [&](const ParsedMarkerCall &call) -> std::optional { + if (call.args.size() != expectedNumArgs) + return std::nullopt; + + std::string replacement; + replacement.reserve(marker.size() + 16); + replacement.append(call.args[0].str()); + replacement.push_back('.'); + replacement.append(memberName.str()); + replacement.push_back('('); + if (expectedNumArgs >= 2) + replacement.append(call.args[1].str()); + if (expectedNumArgs == 3) { + replacement.append(", "); + replacement.append(call.args[2].str()); + } + replacement.push_back(')'); + return replacement; + }); +} + +static void rewriteMarkerCallsToMembers( + std::string &cpp, llvm::ArrayRef rewrites) { + bool changed = true; + while (changed) { + changed = false; + for (const MarkerRewriteSpec &rewrite : rewrites) { + changed |= rewriteMarkerCallToMember(cpp, rewrite.marker, + rewrite.memberName, + rewrite.expectedNumArgs); + } + } +} + +static void rewriteTileGetSetValueMarkers(std::string &cpp) { + static const MarkerRewriteSpec kTileMarkerRewrites[] = { + {"PTOAS__TILE_SET_VALUE", "SetValue", 3}, + {"PTOAS__TILE_GET_VALUE", "GetValue", 2}, + {"PTOAS__TILE_DATA", "data", 1}, + {"PTOAS__TILE_SET_VALIDSHAPE", "SetValidShape", 3}, + }; + rewriteMarkerCallsToMembers(cpp, kTileMarkerRewrites); +} + +static void rewriteAsyncEventMarkers(std::string &cpp) { + static const MarkerRewriteSpec kAsyncEventMarkerRewrites[] = { + {"PTOAS__ASYNC_EVENT_WAIT", "Wait", 2}, + {"PTOAS__ASYNC_EVENT_TEST", "Test", 2}, + }; + rewriteMarkerCallsToMembers(cpp, kAsyncEventMarkerRewrites); +} + +// -------------------------------------------------------------------------- +// EmitC cleanup: drop empty emitc.expression ops. +// +// After FormExpressions + CSE, EmitC expressions can become empty when their +// root op is CSE'd with an equivalent dominating value outside the expression +// region. Such expressions crash mlir::emitc::translateToCpp because +// ExpressionOp::getRootOp() returns nullptr. +// -------------------------------------------------------------------------- +static void dropEmptyEmitCExpressions(Operation *rootOp) { + llvm::SmallVector toErase; + rootOp->walk([&](emitc::ExpressionOp expr) { + if (expr.getRootOp()) + return; + Block *body = expr.getBody(); + if (!body) + return; + auto yield = dyn_cast(body->getTerminator()); + if (!yield || yield.getNumOperands() != 1) + return; + Value yielded = yield.getOperand(0); + expr.getResult().replaceAllUsesWith(yielded); + toErase.push_back(expr); + }); + for (emitc::ExpressionOp expr : llvm::reverse(toErase)) + expr.erase(); +} + +static Attribute getDefaultEmitCVariableInitAttr(OpBuilder &builder, Type type) { + if (auto intTy = dyn_cast(type)) + return builder.getIntegerAttr(intTy, 0); + if (isa(type)) + return builder.getIndexAttr(0); + if (auto floatTy = dyn_cast(type)) + return builder.getFloatAttr(floatTy, 0.0); + if (isa(type)) + return emitc::OpaqueAttr::get(builder.getContext(), ""); + return Attribute{}; +} + +// FormExpressions may inline conditions into emitc.expression, but the C++ +// emitter prints cf.br/cf.cond_br operands by variable name rather than by +// recursively emitting an expression. Materialize such operands so CFG-based +// lowering (e.g. scf.while -> cf.*) stays valid. +static void materializeControlFlowOperands(Operation *rootOp) { + llvm::SmallVector branches; + rootOp->walk([&](Operation *op) { + if (isa(op)) + branches.push_back(op); + }); + + OpBuilder builder(rootOp->getContext()); + for (Operation *op : branches) { + builder.setInsertionPoint(op); + for (OpOperand &operand : op->getOpOperands()) { + Value value = operand.get(); + auto expr = dyn_cast_or_null(value.getDefiningOp()); + if (!expr) + continue; + + Attribute initAttr = + getDefaultEmitCVariableInitAttr(builder, value.getType()); + if (!initAttr) + continue; + + Value tmp = + builder.create(op->getLoc(), value.getType(), + initAttr) + .getResult(); + builder.create(op->getLoc(), tmp, value); + operand.set(tmp); + } + } +} + +static bool rewriteMarkerCallToSubscript(std::string &cpp, llvm::StringRef marker, + unsigned expectedNumArgs, + bool isStore) { + return rewriteMarkerCalls( + cpp, marker, [&](const ParsedMarkerCall &call) -> std::optional { + if (call.args.size() != expectedNumArgs) + return std::nullopt; + if (isStore) { + return (call.args[0] + "[" + call.args[1] + "] = " + call.args[2]) + .str(); + } + return (call.args[0] + "[" + call.args[1] + "]").str(); + }); +} + +static void rewriteMarkerCallsToSubscripts( + std::string &cpp, llvm::ArrayRef rewrites) { + bool changed = true; + while (changed) { + changed = false; + for (const MarkerSubscriptRewriteSpec &rewrite : rewrites) { + changed |= rewriteMarkerCallToSubscript(cpp, rewrite.marker, + rewrite.expectedNumArgs, + rewrite.isStore); + } + } +} + +static void rewritePtrScalarMarkers(std::string &cpp) { + static const MarkerSubscriptRewriteSpec kPtrMarkerRewrites[] = { + {"PTOAS__PTR_LOAD", 2, false}, + {"PTOAS__PTR_STORE", 3, true}, + }; + rewriteMarkerCallsToSubscripts(cpp, kPtrMarkerRewrites); +} + +static void rewriteEventIdArrayMarkers(std::string &cpp) { + static const MarkerSubscriptRewriteSpec kEventIdMarkerRewrites[] = { + {"PTOAS__EVENTID_ARRAY_LOAD", 2, false}, + {"PTOAS__EVENTID_ARRAY_STORE", 3, true}, + }; + rewriteMarkerCallsToSubscripts(cpp, kEventIdMarkerRewrites); +} + +static bool rewriteAddPtrTraceMarkers(std::string &cpp, bool showTrace) { + size_t searchPos = 0; + bool changed = false; + for (auto call = findNextMarkerCall(cpp, "PTOAS__ADDPTR_TRACE", searchPos); + call; call = findNextMarkerCall(cpp, "PTOAS__ADDPTR_TRACE", searchPos)) { + if (call->rparenPos == std::string::npos) { + searchPos = call->markerPos + 1; + continue; + } + if (call->args.size() != 3) { + searchPos = call->rparenPos + 1; + continue; + } + + std::string replacement; + if (showTrace) { + replacement.reserve(64); + replacement.append("/* ADDPTR_TRACE: "); + replacement.append(call->args[0].str()); + replacement.append(" = "); + replacement.append(call->args[1].str()); + replacement.append(" + "); + replacement.append(call->args[2].str()); + replacement.append(" */"); + } + + size_t replaceEnd = call->rparenPos; + if (!showTrace) { + size_t i = call->rparenPos + 1; + while (i < cpp.size() && std::isspace(static_cast(cpp[i]))) + ++i; + if (i < cpp.size() && cpp[i] == ';') + replaceEnd = i; + } + + cpp.replace(call->markerPos, (replaceEnd - call->markerPos) + 1, + replacement); + changed = true; + searchPos = call->markerPos + replacement.size(); + } + return changed; +} + +static bool isGeneratedGlobalTensorDecl(llvm::StringRef trimmed, + llvm::StringRef &decl, + llvm::StringRef &varName) { + if (!trimmed.starts_with("GlobalTensor<") || !trimmed.ends_with(";") || + trimmed.contains('=') || trimmed.contains('(')) { + return false; + } + + decl = trimmed.drop_back().rtrim(); + size_t lastWs = decl.find_last_of(" \t"); + if (lastWs == llvm::StringRef::npos) + return false; + varName = decl.drop_front(lastWs + 1); + if (!varName.starts_with("v") || varName.size() <= 1) + return false; + return llvm::all_of(varName.drop_front(1), + [](char c) { return std::isdigit(c); }); +} + +static void rewriteHoistedGlobalTensorDecls(std::string &cpp) { + // When `declareVariablesAtTop` is enabled, the C++ emitter hoists SSA value + // declarations to the top of the function and emits assignments later. This + // requires the C++ type to be default-constructible. + // + // `GlobalTensor<...>` from pto-isa does NOT have a default constructor, so a + // hoisted declaration like: + // GlobalTensor<...> v42; + // fails to compile. Initialize those hoisted temporaries with a null pointer + // so they are constructible: + // GlobalTensor<...> v42(nullptr); + // + // We keep the assignment later; the null-initialized value is never used. + std::string out; + out.reserve(cpp.size() + 64); + + llvm::StringRef ref(cpp); + while (!ref.empty()) { + auto split = ref.split('\n'); + llvm::StringRef line = split.first; + llvm::StringRef rest = split.second; + + llvm::StringRef trimmed = line.trim(); + bool rewritten = false; + llvm::StringRef decl; + llvm::StringRef varName; + if (isGeneratedGlobalTensorDecl(trimmed, decl, varName)) { + size_t indentLen = line.find_first_not_of(" \t"); + if (indentLen == std::string::npos) + indentLen = 0; + llvm::StringRef indent = line.take_front(indentLen); + + out.append(indent.str()); + out.append(decl.str()); + out.append("(nullptr);"); + rewritten = true; + } + + if (!rewritten) + out.append(line.str()); + if (!rest.empty()) + out.push_back('\n'); + ref = rest; + } + + cpp.swap(out); +} + +namespace { +struct ConstantDeclCandidate { + size_t declLine = 0; + std::string indent; + std::string type; + bool hasInitializer = false; + std::string initializer; + size_t assignmentCount = 0; + size_t assignmentLine = 0; + std::string assignmentRhs; +}; +} // namespace + +static bool isGeneratedValueName(llvm::StringRef name) { + if (!name.consume_front("v") || name.empty()) + return false; + return llvm::all_of(name, [](char c) { return std::isdigit(c); }); +} + +static bool isConstFoldableScalarType(llvm::StringRef type) { + type = type.trim(); + if (type.starts_with("const ") || type.starts_with("constexpr ")) + return false; + return llvm::StringSwitch(type) + .Cases("bool", "float", "double", "half", "bfloat16_t", true) + .Cases("int8_t", "uint8_t", "int16_t", "uint16_t", true) + .Cases("int32_t", "uint32_t", "int64_t", "uint64_t", true) + .Default(false); +} + +static bool isLiteralInitializer(llvm::StringRef rhs) { + rhs = rhs.trim(); + if (rhs.empty()) + return false; + if (rhs == "true" || rhs == "false" || rhs == "nullptr") + return true; + + static const llvm::Regex kIntLiteral( + R"(^[+-]?(0[xX][0-9A-Fa-f]+|[0-9]+)[uUlL]*$)"); + static const llvm::Regex kFloatLiteral( + R"(^[+-]?(([0-9]+\.[0-9]*|\.[0-9]+|[0-9]+)([eE][+-]?[0-9]+)?|[0-9]+[eE][+-]?[0-9]+)[fF]?$)"); + static const llvm::Regex kHexFloatLiteral( + R"(^[+-]?0[xX]([0-9A-Fa-f]+\.[0-9A-Fa-f]*|[0-9A-Fa-f]+|\.[0-9A-Fa-f]+)[pP][+-]?[0-9]+[fF]?$)"); + static const llvm::Regex kSpecialFloatLiteral( + R"(^[+-]?(nan|inf)[fF]?$)"); + + return kIntLiteral.match(rhs) || kFloatLiteral.match(rhs) || + kHexFloatLiteral.match(rhs) || kSpecialFloatLiteral.match(rhs); +} + +static std::string normalizeConstInitializer(llvm::StringRef type, + llvm::StringRef rhs) { + type = type.trim(); + rhs = rhs.trim(); + if (type == "bool") { + if (rhs == "0" || rhs == "false") + return "false"; + if (rhs == "1" || rhs == "-1" || rhs == "true") + return "true"; + } + return rhs.str(); +} + +static bool parseConstantDeclarationLine(llvm::StringRef line, + ConstantDeclCandidate &candidate, + std::string &valueName) { + llvm::StringRef trimmed = line.trim(); + if (trimmed.empty() || trimmed.starts_with("#") || trimmed.starts_with("//") || + !trimmed.ends_with(";")) + return false; + + llvm::StringRef body = trimmed.drop_back().rtrim(); + if (body.starts_with("return") || body.starts_with("goto ") || + body.starts_with("if ") || body.starts_with("if(") || + body.starts_with("switch ") || body.starts_with("switch(") || + body.starts_with("for ") || body.starts_with("for(") || + body.starts_with("while ") || body.starts_with("while(") || + body.starts_with("case ") || body == "default") + return false; + + llvm::StringRef lhs = body; + llvm::StringRef rhs; + if (size_t eqPos = body.find('='); eqPos != llvm::StringRef::npos) { + lhs = body.take_front(eqPos).rtrim(); + rhs = body.drop_front(eqPos + 1).trim(); + } + + size_t lastWs = lhs.find_last_of(" \t"); + if (lastWs == llvm::StringRef::npos) + return false; + + llvm::StringRef type = lhs.take_front(lastWs).rtrim(); + llvm::StringRef name = lhs.drop_front(lastWs + 1).trim(); + if (!isGeneratedValueName(name) || !isConstFoldableScalarType(type)) + return false; + + size_t indentLen = line.find_first_not_of(" \t"); + if (indentLen == llvm::StringRef::npos) + indentLen = 0; + candidate.indent = line.take_front(indentLen).str(); + candidate.type = type.str(); + valueName = name.str(); + + if (!rhs.empty()) { + if (!isLiteralInitializer(rhs)) + return false; + candidate.hasInitializer = true; + candidate.initializer = normalizeConstInitializer(type, rhs); + } + + return true; +} + +static bool parseGeneratedValueAssignment(llvm::StringRef line, + llvm::StringRef &valueName, + llvm::StringRef &rhs) { + llvm::StringRef trimmed = line.trim(); + if (trimmed.empty() || trimmed.starts_with("#") || trimmed.starts_with("//") || + !trimmed.ends_with(";")) + return false; + + llvm::StringRef body = trimmed.drop_back().rtrim(); + size_t eqPos = body.find('='); + if (eqPos == llvm::StringRef::npos) + return false; + + llvm::StringRef lhs = body.take_front(eqPos).rtrim(); + rhs = body.drop_front(eqPos + 1).trim(); + if (!isGeneratedValueName(lhs)) + return false; + valueName = lhs; + return true; +} + +static void rewriteScalarConstantDecls(std::string &cpp) { + llvm::SmallVector lines; + for (llvm::StringRef ref(cpp); !ref.empty(); ref = ref.split('\n').second) { + auto split = ref.split('\n'); + lines.push_back(split.first.str()); + } + + llvm::SmallVector eraseLine(lines.size(), false); + auto rewriteSegment = [&](size_t beginLine, size_t endLine) { + llvm::StringMap candidates; + + for (size_t i = beginLine; i <= endLine; ++i) { + ConstantDeclCandidate candidate; + std::string valueName; + if (parseConstantDeclarationLine(lines[i], candidate, valueName)) { + candidate.declLine = i; + candidates[valueName] = std::move(candidate); + continue; + } + + llvm::StringRef assignedName; + llvm::StringRef rhs; + if (!parseGeneratedValueAssignment(lines[i], assignedName, rhs)) + continue; + + auto it = candidates.find(assignedName); + if (it == candidates.end()) + continue; + + ConstantDeclCandidate &info = it->second; + ++info.assignmentCount; + info.assignmentLine = i; + info.assignmentRhs = rhs.str(); + } + + for (auto &entry : candidates) { + llvm::StringRef valueName = entry.getKey(); + ConstantDeclCandidate &info = entry.getValue(); + + std::string initializer; + if (info.hasInitializer) { + if (info.assignmentCount != 0) + continue; + initializer = info.initializer; + } else { + if (info.assignmentCount != 1) + continue; + if (!isLiteralInitializer(info.assignmentRhs)) + continue; + initializer = normalizeConstInitializer( + info.type, llvm::StringRef(info.assignmentRhs)); + eraseLine[info.assignmentLine] = true; + } + + lines[info.declLine] = (info.indent + "const " + info.type + " " + + valueName.str() + " = " + initializer + ";"); + } + }; + + int braceDepth = 0; + size_t segmentStart = 0; + for (size_t i = 0; i < lines.size(); ++i) { + int depthBefore = braceDepth; + for (char c : lines[i]) { + if (c == '{') + ++braceDepth; + else if (c == '}') + --braceDepth; + } + + if (depthBefore == 0 && braceDepth > 0) + segmentStart = i; + if (depthBefore > 0 && braceDepth == 0) + rewriteSegment(segmentStart, i); + } + + std::string out; + out.reserve(cpp.size()); + for (size_t i = 0; i < lines.size(); ++i) { + if (eraseLine[i]) + continue; + out.append(lines[i]); + if (i + 1 != lines.size()) + out.push_back('\n'); + } + cpp.swap(out); +} + +static bool shouldDeclareVariablesAtTop(ModuleOp module) { + auto hasMultiBlockFunc = [](auto func) { return func.getBlocks().size() > 1; }; + return llvm::any_of(module.getOps(), hasMultiBlockFunc) || + llvm::any_of(module.getOps(), hasMultiBlockFunc); +} + +static void prepareVPTOForEmission(PassManager &pm) { + auto &kernelModulePM = pm.nest(); + kernelModulePM.addPass(createCanonicalizerPass()); + kernelModulePM.addPass(createCSEPass()); + kernelModulePM.addPass(pto::createVPTOPtrNormalizePass()); + kernelModulePM.addPass(pto::createVPTOPtrCastCleanupPass()); + kernelModulePM.addPass(createReconcileUnrealizedCastsPass()); + kernelModulePM.addNestedPass( + createVPTOExpandWrapperOpsPass()); + kernelModulePM.addPass(createCSEPass()); + kernelModulePM.addNestedPass( + pto::createPTOInferVPTOVecScopePass()); + kernelModulePM.addPass(createCanonicalizerPass()); + kernelModulePM.addPass(createCSEPass()); + kernelModulePM.addPass(pto::createPTOValidateVPTOEmissionIRPass()); +} + +static void lowerPTOToVPTOBackend(PassManager &pm, int argc, char **argv) { + // TileOp Expand path: + // 1. ExpandTileOp: instantiate TileLang DSL templates, replace tile ops + // with func.call to template functions (tile_buf params) + // 2. InlineLibCall: inline template function bodies + // 3. FoldTileBufIntrinsics: fold tile_buf_addr / tile_valid_rows / + // tile_valid_cols to concrete memref/constant values + auto &kernelModulePM = pm.nest(); + pto::ExpandTileOpOptions expandOpts = resolveExpandTileOpOptions(argc, argv); + kernelModulePM.addPass(pto::createExpandTileOpPass(expandOpts)); + + kernelModulePM.addPass(pto::createPTOInlineLibCallPass()); + kernelModulePM.addNestedPass( + pto::createFoldTileBufIntrinsicsPass()); + // FoldTileBufIntrinsics materializes many constant branch conditions. + // Clean them up immediately on the TileOp expansion path before the + // authoring-stage VPTO verifier and let the existing CSE passes remove the + // resulting dead values later in the pipeline. + kernelModulePM.addPass(mlir::createSCCPPass()); + kernelModulePM.addPass(mlir::createCanonicalizerPass()); +} + +static void inlineTilelangHelpersOnVPTOInput(PassManager &pm) { + auto &kernelModulePM = pm.nest(); + kernelModulePM.addPass(pto::createPTOInlineLibCallPass()); + kernelModulePM.addPass(mlir::createSCCPPass()); + kernelModulePM.addPass(mlir::createCanonicalizerPass()); +} + +static pto::VPTOEmissionOptions buildVPTOEmissionOptions() { + pto::VPTOEmissionOptions options; + options.dumpVPTOIR = false; + options.targetTriple = "hiipu64-hisilicon-cce"; + return options; +} + +static LogicalResult runVPTOBackendPipeline(ModuleOp module, int argc, + char **argv, + bool hasTileOpsToExpand, + bool hasTilelangHelpers); + +static LogicalResult collectVPTOBackendResult(ModuleOp module, + PTOASCompileResult &result, + llvm::raw_ostream &diagOS, + bool emitHostStub) { + if (emitVPTO) { + result.kind = PTOASCompileResultKind::Text; + llvm::raw_string_ostream os(result.textOutput); + module.print(os); + os << "\n"; + os.flush(); + return success(); + } + + pto::VPTOEmissionOptions options = buildVPTOEmissionOptions(); + result.kind = PTOASCompileResultKind::VPTOObject; + if (emitHostStub) { + if (failed(pto::emitVPTOHostStubSource(module, result.stubSource, diagOS))) { + diagOS << "Error: Failed to emit VPTO host stub source.\n"; + return failure(); + } + } + + pto::EmittedLLVMModule cubeModule; + pto::EmittedLLVMModule vectorModule; + if (failed(pto::lowerVPTOModuleToLLVMModules(module, options, cubeModule, + vectorModule, diagOS))) { + diagOS << "Error: Failed to lower VPTO to LLVM modules.\n"; + return failure(); + } + + if (vectorModule.module) { + PTOASBackendObjectInput object; + object.kind = PTOASBackendObjectKind::VPTOVectorDeviceObject; + object.llvmModule = std::move(vectorModule); + result.backendObjects.push_back(std::move(object)); + } + if (cubeModule.module) { + PTOASBackendObjectInput object; + object.kind = PTOASBackendObjectKind::VPTOCubeDeviceObject; + object.llvmModule = std::move(cubeModule); + result.backendObjects.push_back(std::move(object)); + } + return success(); +} + +static LogicalResult runVPTOBackendPipeline(ModuleOp module, int argc, + char **argv, + bool hasTileOpsToExpand, + bool hasTilelangHelpers) { + PassManager pm(module.getContext()); + pm.enableVerifier(); + pm.addPass(pto::createVPTOSplitCVModulePass()); + pm.addPass(pto::createVPTONormalizeContainerPass()); + if (!hasTileOpsToExpand && hasTilelangHelpers) + inlineTilelangHelpersOnVPTOInput(pm); + if (hasTileOpsToExpand) + lowerPTOToVPTOBackend(pm, argc, argv); + prepareVPTOForEmission(pm); + if (failed(applyConfiguredPassManagerCLOptions( + pm, "VPTO unified emission pipeline"))) + return failure(); + if (failed(pm.run(module))) { + llvm::errs() << "Error: VPTO emission pipeline failed.\n"; + return failure(); + } + if (failed(applyVPTOPublicABIExportNames(module))) + return failure(); + return success(); +} + +LogicalResult mlir::pto::compilePTOASModule(ModuleOp module, + llvm::StringRef arch, + bool cliBackendSpecified, int argc, + char **argv, + PTOASCompileResult &result, + bool emitVPTOHostStub) { + result = PTOASCompileResult{}; + PTOASBackend effectiveBackend = PTOASBackend::EmitC; + if (!parseBackend(ptoBackend, effectiveBackend)) { + llvm::errs() << "Error: invalid --pto-backend='" << ptoBackend + << "'. Expected 'emitc' or 'vpto'.\n"; + return failure(); + } + + bool mixedBackendMode = false; + if (failed(resolveModuleBackend(module, cliBackendSpecified, + mixedBackendMode, effectiveBackend))) + return failure(); + if (failed(verifyPTOASBackendInput(module, mixedBackendMode))) + return failure(); + + CompilationMode compilationMode = effectiveBackend == PTOASBackend::VPTO + ? CompilationMode::VPTO + : CompilationMode::EmitC; + + if (effectiveBackend != PTOASBackend::VPTO && + (emitVPTO || ptoPrintSeamIR || !ptoSeamIRFile.empty())) { + llvm::errs() << "Error: VPTO-specific flags require " + "--pto-backend=vpto or pto.backend = \"vpto\".\n"; + return failure(); + } + + PTOBuildLevel effectiveLevel = defaultBuildLevel(); + if (!parseBuildLevel(ptoBuildLevel, effectiveLevel)) { + llvm::errs() << "Error: invalid --pto-level='" << ptoBuildLevel + << "'. Expected 'level1', 'level2', or 'level3'.\n"; + return failure(); + } + + bool invalidAutoSyncTailHint = false; + module.walk([&](mlir::func::FuncOp func) { + auto hintAttr = + func->getAttrOfType("pto.auto_sync_tail_hint"); + if (!hintAttr) + return; + + std::string normalizedHint; + if (!parseAutoSyncTailHint(hintAttr.getValue(), normalizedHint)) { + func.emitError("invalid pto.auto_sync_tail_hint '") + << hintAttr.getValue() + << "'. Expected 'barrier-all' (or 'default') or " + "'mte3-to-s-event0'."; + invalidAutoSyncTailHint = true; + return; + } + func->setAttr("pto.auto_sync_tail_hint", + mlir::StringAttr::get(module.getContext(), normalizedHint)); + }); + if (invalidAutoSyncTailHint) + return failure(); + + bool hasTAssign = false; + module.walk([&](pto::TAssignOp) { hasTAssign = true; }); + + if (hasTAssign && effectiveLevel != PTOBuildLevel::Level3) { + llvm::errs() << "Error: pto.tassign is only supported when " + "--pto-level=level3.\n"; + return failure(); + } + + if (hasTAssign && enableInsertSync) { + llvm::errs() << "Error: pto.tassign requires --enable-insert-sync to be " + "disabled.\n"; + return failure(); + } + + int enabledAutoSyncModes = + (enableInsertSync ? 1 : 0) + (enableInjectBarrierAllSync ? 1 : 0) + + (enableGraphSyncSolver ? 1 : 0); + if (enabledAutoSyncModes > 1) { + llvm::errs() << "Error: --enable-insert-sync, " + "--enable-inject-barrier-all-sync, and " + "--enable-graph-sync-solver are mutually exclusive.\n"; + return failure(); + } + if (hasTAssign && enableInjectBarrierAllSync) { + llvm::errs() << "Error: pto.tassign requires " + "--enable-inject-barrier-all-sync to be disabled.\n"; + return failure(); + } + if (hasTAssign && enableGraphSyncSolver) { + llvm::errs() << "Error: pto.tassign requires --enable-graph-sync-solver " + "to be disabled.\n"; + return failure(); + } + + if (effectiveLevel == PTOBuildLevel::Level3) { + bool missing = false; + module.walk([&](pto::AllocTileOp op) { + if (!op.getAddr()) { + op.emitError("requires 'addr' operand when --pto-level=level3"); + missing = true; + } + }); + if (missing) + return failure(); + } else { + bool hasAddr = false; + module.walk([&](pto::AllocTileOp op) { + if (op.getAddr()) { + op.emitError( + "unexpected 'addr' operand: only supported when --pto-level=level3"); + hasAddr = true; + } + }); + if (hasAddr) + return failure(); + } + + const bool hasTileOpsToExpand = hasUnexpandedTileOps(module); + const bool hasTilelangHelpers = hasTilelangInlineHelpers(module); + + if (compilationMode == CompilationMode::VPTO && !hasTileOpsToExpand) { + if (ptoPrintSeamIR || !ptoSeamIRFile.empty()) { + llvm::errs() << "Error: shared pre-backend seam IR is unavailable when " + "skipping the shared PTO-to-VPTO lowering pipeline.\n"; + return failure(); + } + if (failed(runVPTOBackendPipeline(module, argc, argv, hasTileOpsToExpand, + hasTilelangHelpers))) + return failure(); + return collectVPTOBackendResult(module, result, llvm::errs(), + emitVPTOHostStub); + } + + // Main PassManager + PassManager pm(module.getContext()); + + if (failed(applyPassManagerCLOptions(pm))) + return failure(); + + pm.addNestedPass( + pto::createPTOAssignDefaultFrontendPipeIdPass()); + pm.addNestedPass( + pto::createPTOLowerFrontendPipeOpsPass()); + // pm.addNestedPass(pto::createPTOVerifyTFreePass()); + pm.addPass(pto::createPTOInferValidatePipeInitPass()); + pm.addNestedPass(pto::createLoweringSyncToPipePass()); + + if (!disableInferLayout) + pm.addNestedPass(pto::createInferPTOLayoutPass()); + pm.addNestedPass(pto::createPTOA5NormalizeTMovPass()); + pm.addPass(pto::createPTOViewToMemrefPass()); + + if (effectiveLevel != PTOBuildLevel::Level3) { + PlanMemoryOptions planMemoryOption; + planMemoryOption.memMode = MemPlanMode::LOCAL_MEM_PLAN; + planMemoryOption.enableGlobalReuse = false; + planMemoryOption.enablePrintMemoryAllocatedSize = false; + pm.addPass(pto::createPlanMemoryPass(planMemoryOption)); + } + pm.addPass(pto::createPTOResolveReservedBuffersPass()); + + // Conditionally add one automatic synchronization mode. Barrier-all is a + // conservative standalone pass; InsertSync and GraphSyncSolver are set/wait + // solvers. + if (enableInsertSync) + pm.addNestedPass(pto::createPTOInsertSyncPass()); + else if (enableInjectBarrierAllSync) + pm.addNestedPass( + pto::createPTOInjectBarrierAllSyncPass()); + else if (enableGraphSyncSolver) { + PTOGraphSyncSolverOptions graphSyncOpts; + graphSyncOpts.eventIdNumMax = graphSyncSolverEventIdMax; + pm.addNestedPass( + pto::createPTOGraphSyncSolverPass(graphSyncOpts)); + } + + if (emitMlirIR) { + if (failed(pm.run(module))) { + llvm::errs() << "Error: Pass execution failed.\n"; + return failure(); + } + result.kind = PTOASCompileResultKind::Text; + llvm::raw_string_ostream os(result.textOutput); + module.print(os); + os.flush(); + return success(); + } + + // Reintroduce tile-native handles once on the shared mainline so both + // backends consume the same post-planning seam IR. + pm.addPass(pto::createPTOMaterializeTileHandlesPass()); + pm.addPass(createCSEPass()); + if (failed(applyConfiguredPassManagerCLOptions(pm, "main PTOAS pipeline"))) + return failure(); + + module->setAttr("pto.target_arch", + mlir::StringAttr::get(module.getContext(), arch)); + + if (compilationMode == CompilationMode::VPTO) { + if (failed(pm.run(module))) { + llvm::errs() << "Error: Pass execution failed.\n"; + return failure(); + } + + if (ptoPrintSeamIR) { + module.print(llvm::errs()); + llvm::errs() << "\n"; + } + if (failed(emitSharedPreBackendSeamIR(module, ptoSeamIRFile))) + return failure(); + + if (failed(runVPTOBackendPipeline(module, argc, argv, hasTileOpsToExpand, + hasTilelangHelpers))) + return failure(); + return collectVPTOBackendResult(module, result, llvm::errs(), + emitVPTOHostStub); + } + + if (arch == "a3") { + pm.addPass(pto::createEmitPTOManualPass(pto::PTOArch::A3)); + } else { + pm.addPass(pto::createEmitPTOManualPass(pto::PTOArch::A5)); + } + pm.addPass(emitc::createFormExpressionsPass()); + pm.addPass(mlir::createCSEPass()); + + if (failed(pm.run(module))) { + llvm::errs() << "Error: Pass execution failed.\n"; + return failure(); + } + + std::string cppOutput; + dropEmptyEmitCExpressions(module); + materializeControlFlowOperands(module); + if (failed(reorderEmitCFunctions(module))) { + llvm::errs() << "Error: Failed to order emitted functions for C++ " + "emission.\n"; + return failure(); + } + + llvm::raw_string_ostream cppOS(cppOutput); + bool declareVariablesAtTop = shouldDeclareVariablesAtTop(module); + if (failed(emitc::translateToCpp( + module, cppOS, /*declareVariablesAtTop=*/declareVariablesAtTop))) { + llvm::errs() << "Error: Failed to emit C++.\n"; + return failure(); + } + cppOS.flush(); + rewriteTileGetSetValueMarkers(cppOutput); + rewriteAsyncEventMarkers(cppOutput); + rewritePtrScalarMarkers(cppOutput); + rewriteEventIdArrayMarkers(cppOutput); + rewriteAddPtrTraceMarkers(cppOutput, emitAddPtrTrace); + rewriteScalarConstantDecls(cppOutput); + rewriteHoistedGlobalTensorDecls(cppOutput); + + result.kind = PTOASCompileResultKind::Text; + result.textOutput = std::move(cppOutput); + return success(); } diff --git a/tools/ptoas/ptoas.h b/tools/ptoas/ptoas.h new file mode 100644 index 0000000000..4bd0fae797 --- /dev/null +++ b/tools/ptoas/ptoas.h @@ -0,0 +1,68 @@ +// 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. + +#ifndef PTOAS_H +#define PTOAS_H + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Support/LogicalResult.h" +#include "PTO/Transforms/VPTOLLVMEmitter.h" +#include "llvm/ADT/StringRef.h" +#include +#include +#include +#include + +namespace mlir::pto { + +enum class PTOASBackend { + EmitC, + VPTO, +}; + +enum class PTOASCompileResultKind { + Text, + VPTOObject, + MixedObject, +}; + +enum class PTOASBackendObjectKind { + EmitCFatobj, + VPTOVectorDeviceObject, + VPTOCubeDeviceObject, +}; + +struct PTOASBackendObjectInput { + PTOASBackendObjectKind kind = PTOASBackendObjectKind::EmitCFatobj; + std::string cppSource; + EmittedLLVMModule llvmModule; +}; + +struct PTOASCompileResult { + PTOASCompileResultKind kind = PTOASCompileResultKind::Text; + std::string textOutput; + std::string stubSource; + std::vector backendObjects; +}; + +llvm::StringRef getPTOASTargetArchOption(); +LogicalResult getPTOASCommandLineBackend(PTOASBackend &backend); +bool isPTOASDebugIROutputRequested(); +std::string normalizePTOASArch(llvm::StringRef archValue); +bool isSupportedPTOASArch(llvm::StringRef archValue); +std::optional detectPTOASTextualModuleArch(llvm::StringRef text); + +LogicalResult compilePTOASModule(ModuleOp module, llvm::StringRef arch, + bool cliBackendSpecified, int argc, + char **argv, + PTOASCompileResult &result, + bool emitVPTOHostStub = true); + +} // namespace mlir::pto + +#endif