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/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/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/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_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..a2c929e75a --- /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 requires an explicit file path passed with -o 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..9ea92f43c5 --- /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 %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 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_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..9f8892de01 --- /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: 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 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/tools/ptoas/VPTOFatobjEmission.h b/test/vpto/cases/micro-op/backend/mixed-external-vadd/launch.cpp similarity index 53% rename from tools/ptoas/VPTOFatobjEmission.h rename to test/vpto/cases/micro-op/backend/mixed-external-vadd/launch.cpp index fd161ba608..97c1e01cda 100644 --- a/tools/ptoas/VPTOFatobjEmission.h +++ b/test/vpto/cases/micro-op/backend/mixed-external-vadd/launch.cpp @@ -6,26 +6,20 @@ // 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; -} - -namespace mlir::pto { +#ifndef __VEC_SCOPE__ +#define __VEC_SCOPE__ +#endif -LogicalResult emitVPTOFatobj(llvm::Module *cubeModule, - llvm::Module *vectorModule, - llvm::StringRef stubSource, - llvm::ToolOutputFile &outputFile, - llvm::raw_ostream &diagOS); +#include +#ifndef __CPU_SIM +#include "acl/acl.h" +#endif -} // namespace mlir::pto +extern "C" __global__ [aicore] void mixed_external_vadd_kernel( + __gm__ float *lhs, __gm__ float *rhs, __gm__ float *out); -#endif +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..ced12af999 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 + driver.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/VPTOHostStubEmission.cpp b/tools/ptoas/VPTOHostStubEmission.cpp index ae56d85f8a..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; @@ -74,35 +70,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() || !pto::isPTOKernelFunction(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 +108,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/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 caa32e5b36..4f5cdd7506 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -6,11 +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. +#include "ptoas.h" + #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" @@ -25,7 +26,6 @@ #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" @@ -34,7 +34,6 @@ #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" @@ -49,8 +48,6 @@ #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 @@ -59,10 +56,6 @@ #include #include -extern "C" { -extern char **environ; -} - using namespace mlir; using namespace pto; @@ -70,12 +63,6 @@ using namespace pto; #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); @@ -95,7 +82,8 @@ static std::string joinPath(llvm::StringRef lhs, llvm::StringRef rhs) { } static std::string detectInstalledTilelangPath(const char *argv0) { - std::string exePath = llvm::sys::fs::getMainExecutable(argv0, (void *)&main); + std::string exePath = + llvm::sys::fs::getMainExecutable(argv0, (void *)&compilePTOASModule); if (exePath.empty()) return {}; @@ -108,7 +96,8 @@ static std::string detectInstalledTilelangPath(const char *argv0) { } static std::string detectInstalledTilelangPkgPath(const char *argv0) { - std::string exePath = llvm::sys::fs::getMainExecutable(argv0, (void *)&main); + std::string exePath = + llvm::sys::fs::getMainExecutable(argv0, (void *)&compilePTOASModule); if (exePath.empty()) return {}; @@ -253,18 +242,6 @@ static LogicalResult reorderEmitCFunctions(ModuleOp module) { 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)); @@ -434,15 +411,39 @@ enum class PTOBuildLevel { Level3, }; -enum class PTOBackend { - EmitC, - VPTO, -}; - 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) @@ -484,21 +485,180 @@ static bool parseAutoSyncTailHint(llvm::StringRef hintStr, std::string &normaliz return false; } -static bool parseBackend(llvm::StringRef backendStr, PTOBackend &out) { +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 = PTOBackend::EmitC; + out = PTOASBackend::EmitC; return true; } if (s == "vpto") { - out = PTOBackend::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()) @@ -1233,46 +1393,61 @@ static pto::VPTOEmissionOptions buildVPTOEmissionOptions() { return options; } -static int emitVPTOBackendResult(ModuleOp module, - llvm::ToolOutputFile &outputFile) { +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) { - module.print(outputFile.os()); - outputFile.os() << "\n"; - outputFile.keep(); - return 0; + result.kind = PTOASCompileResultKind::Text; + llvm::raw_string_ostream os(result.textOutput); + module.print(os); + os << "\n"; + os.flush(); + return success(); } 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; + 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, llvm::errs()))) { - llvm::errs() << "Error: Failed to lower VPTO to LLVM modules.\n"; - return 1; + if (failed(pto::lowerVPTOModuleToLLVMModules(module, options, cubeModule, + vectorModule, diagOS))) { + diagOS << "Error: Failed to lower VPTO to LLVM modules.\n"; + return failure(); } - 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; + if (vectorModule.module) { + PTOASBackendObjectInput object; + object.kind = PTOASBackendObjectKind::VPTOVectorDeviceObject; + object.llvmModule = std::move(vectorModule); + result.backendObjects.push_back(std::move(object)); } - outputFile.keep(); - return 0; + 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(OwningOpRef &module, - int argc, char **argv, +static LogicalResult runVPTOBackendPipeline(ModuleOp module, int argc, + char **argv, bool hasTileOpsToExpand, bool hasTilelangHelpers) { - PassManager pm(module->getContext()); + PassManager pm(module.getContext()); pm.enableVerifier(); pm.addPass(pto::createVPTOSplitCVModulePass()); pm.addPass(pto::createVPTONormalizeContainerPass()); @@ -1284,169 +1459,56 @@ static LogicalResult runVPTOBackendPipeline(OwningOpRef &module, if (failed(applyConfiguredPassManagerCLOptions( pm, "VPTO unified emission pipeline"))) return failure(); - if (failed(pm.run(module.get()))) { + if (failed(pm.run(module))) { llvm::errs() << "Error: VPTO emission pipeline failed.\n"; return failure(); } + if (failed(applyVPTOPublicABIExportNames(module))) + return failure(); return success(); } -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; +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 1; + return failure(); } - 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]); - }; + bool mixedBackendMode = false; + if (failed(resolveModuleBackend(module, cliBackendSpecified, + mixedBackendMode, effectiveBackend))) + return failure(); + if (failed(verifyPTOASBackendInput(module, mixedBackendMode))) + return failure(); - 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; - } - } + CompilationMode compilationMode = effectiveBackend == PTOASBackend::VPTO + ? CompilationMode::VPTO + : CompilationMode::EmitC; - // 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)); + 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 1; + return failure(); } bool invalidAutoSyncTailHint = false; - module->walk([&](mlir::func::FuncOp func) { + module.walk([&](mlir::func::FuncOp func) { auto hintAttr = func->getAttrOfType("pto.auto_sync_tail_hint"); if (!hintAttr) @@ -1462,24 +1524,24 @@ int main(int argc, char **argv) { return; } func->setAttr("pto.auto_sync_tail_hint", - mlir::StringAttr::get(&context, normalizedHint)); + mlir::StringAttr::get(module.getContext(), normalizedHint)); }); if (invalidAutoSyncTailHint) - return 1; + return failure(); bool hasTAssign = false; - module->walk([&](pto::TAssignOp) { hasTAssign = true; }); + 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; + return failure(); } if (hasTAssign && enableInsertSync) { llvm::errs() << "Error: pto.tassign requires --enable-insert-sync to be " "disabled.\n"; - return 1; + return failure(); } int enabledAutoSyncModes = @@ -1489,32 +1551,32 @@ int main(int argc, char **argv) { llvm::errs() << "Error: --enable-insert-sync, " "--enable-inject-barrier-all-sync, and " "--enable-graph-sync-solver are mutually exclusive.\n"; - return 1; + return failure(); } if (hasTAssign && enableInjectBarrierAllSync) { llvm::errs() << "Error: pto.tassign requires " "--enable-inject-barrier-all-sync to be disabled.\n"; - return 1; + return failure(); } if (hasTAssign && enableGraphSyncSolver) { llvm::errs() << "Error: pto.tassign requires --enable-graph-sync-solver " "to be disabled.\n"; - return 1; + return failure(); } if (effectiveLevel == PTOBuildLevel::Level3) { bool missing = false; - module->walk([&](pto::AllocTileOp op) { + module.walk([&](pto::AllocTileOp op) { if (!op.getAddr()) { op.emitError("requires 'addr' operand when --pto-level=level3"); missing = true; } }); if (missing) - return 1; + return failure(); } else { bool hasAddr = false; - module->walk([&](pto::AllocTileOp op) { + module.walk([&](pto::AllocTileOp op) { if (op.getAddr()) { op.emitError( "unexpected 'addr' operand: only supported when --pto-level=level3"); @@ -1522,46 +1584,39 @@ int main(int argc, char **argv) { } }); if (hasAddr) - return 1; + return failure(); } - // [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); - const bool hasTileOpsToExpand = hasUnexpandedTileOps(*module); - const bool hasTilelangHelpers = hasTilelangInlineHelpers(*module); - - if (effectiveBackend == PTOBackend::VPTO && !hasTileOpsToExpand) { + 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 1; + return failure(); } if (failed(runVPTOBackendPipeline(module, argc, argv, hasTileOpsToExpand, hasTilelangHelpers))) - return 1; - return emitVPTOBackendResult(module.get(), outputFile); + return failure(); + return collectVPTOBackendResult(module, result, llvm::errs(), + emitVPTOHostStub); } // Main PassManager - PassManager pm(&context); + PassManager pm(module.getContext()); if (failed(applyPassManagerCLOptions(pm))) - return 1; - + return failure(); + pm.addNestedPass( pto::createPTOAssignDefaultFrontendPipeIdPass()); pm.addNestedPass( pto::createPTOLowerFrontendPipeOpsPass()); - //pm.addNestedPass(pto::createPTOVerifyTFreePass()); + // pm.addNestedPass(pto::createPTOVerifyTFreePass()); pm.addPass(pto::createPTOInferValidatePipeInitPass()); pm.addNestedPass(pto::createLoweringSyncToPipePass()); - + if (!disableInferLayout) pm.addNestedPass(pto::createInferPTOLayoutPass()); pm.addNestedPass(pto::createPTOA5NormalizeTMovPass()); @@ -1591,16 +1646,16 @@ int main(int argc, char **argv) { pto::createPTOGraphSyncSolverPass(graphSyncOpts)); } - llvm::raw_ostream *outputOS = &outputFile.os(); - if (emitMlirIR) { - if (failed(pm.run(*module))) { + if (failed(pm.run(module))) { llvm::errs() << "Error: Pass execution failed.\n"; - return 1; + return failure(); } - module->print(*outputOS); - outputFile.keep(); - return 0; + 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 @@ -1608,28 +1663,29 @@ int main(int argc, char **argv) { pm.addPass(pto::createPTOMaterializeTileHandlesPass()); pm.addPass(createCSEPass()); if (failed(applyConfiguredPassManagerCLOptions(pm, "main PTOAS pipeline"))) - return 1; + return failure(); - module->getOperation()->setAttr("pto.target_arch", - mlir::StringAttr::get(&context, arch)); + module->setAttr("pto.target_arch", + mlir::StringAttr::get(module.getContext(), arch)); - if (effectiveBackend == PTOBackend::VPTO) { - if (failed(pm.run(*module))) { + if (compilationMode == CompilationMode::VPTO) { + if (failed(pm.run(module))) { llvm::errs() << "Error: Pass execution failed.\n"; - return 1; + return failure(); } if (ptoPrintSeamIR) { - module->print(llvm::errs()); + module.print(llvm::errs()); llvm::errs() << "\n"; } - if (failed(emitSharedPreBackendSeamIR(*module, ptoSeamIRFile))) - return 1; + if (failed(emitSharedPreBackendSeamIR(module, ptoSeamIRFile))) + return failure(); if (failed(runVPTOBackendPipeline(module, argc, argv, hasTileOpsToExpand, hasTilelangHelpers))) - return 1; - return emitVPTOBackendResult(module.get(), outputFile); + return failure(); + return collectVPTOBackendResult(module, result, llvm::errs(), + emitVPTOHostStub); } if (arch == "a3") { @@ -1640,29 +1696,26 @@ int main(int argc, char **argv) { pm.addPass(emitc::createFormExpressionsPass()); pm.addPass(mlir::createCSEPass()); - if (failed(pm.run(*module))) { + if (failed(pm.run(module))) { llvm::errs() << "Error: Pass execution failed.\n"; - return 1; + return failure(); } - 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; + 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(); } - // 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))) { + bool declareVariablesAtTop = shouldDeclareVariablesAtTop(module); + if (failed(emitc::translateToCpp( + module, cppOS, /*declareVariablesAtTop=*/declareVariablesAtTop))) { llvm::errs() << "Error: Failed to emit C++.\n"; - return 1; + return failure(); } cppOS.flush(); rewriteTileGetSetValueMarkers(cppOutput); @@ -1672,11 +1725,8 @@ int main(int argc, char **argv) { rewriteAddPtrTraceMarkers(cppOutput, emitAddPtrTrace); rewriteScalarConstantDecls(cppOutput); rewriteHoistedGlobalTensorDecls(cppOutput); - - *outputOS << cppOutput; - outputOS->flush(); - outputFile.keep(); // Success, keep the file - - return 0; + 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