From 5cd46e09db85c544b61f9f308d0fa7f845a496dc Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Fri, 29 May 2026 13:05:12 +0800 Subject: [PATCH] feat: add a driver to manage the compile and link jobs --- .../module-backend-attr-and-mixed-fatobj.md | 442 +++++++ include/PTO/IR/PTO.h | 8 + include/PTO/Transforms/VPTOLLVMEmitter.h | 5 + lib/PTO/IR/PTO.cpp | 9 + lib/PTO/Transforms/PTOToEmitC.cpp | 5 +- lib/PTO/Transforms/PTOViewToMemref.cpp | 11 +- lib/PTO/Transforms/VPTOLLVMEmitter.cpp | 9 +- lib/PTO/Transforms/VPTOSplitCVModule.cpp | 14 +- test/lit/vpto/backend_attr_invalid.pto | 21 + test/lit/vpto/backend_attr_selects_vpto.pto | 26 + ...backend_child_attr_inherits_outer_arch.pto | 29 + .../vpto/backend_child_attr_selects_vpto.pto | 29 + test/lit/vpto/backend_cli_override.pto | 23 + .../lit/vpto/backend_cli_override_to_vpto.pto | 24 + .../backend_mixed_child_attrs_rejected.pto | 28 + ...d_mixed_child_missing_backend_defaults.pto | 28 + ...backend_mixed_external_import_resolves.pto | 35 + .../backend_mixed_requires_output_file.pto | 28 + .../vpto/backend_public_abi_export_suffix.pto | 38 + ...vpto_child_missing_kernel_kind_invalid.pto | 19 + test/lit/vpto/kernel_attr_vpto_llvm.pto | 17 + test/lit/vpto/legacy_aicore_kernel_attr.pto | 4 +- test/lit/vpto/section_sugar_mixed.pto | 3 +- test/lit/vpto/section_sugar_multi_func.pto | 3 +- .../backend/mixed-external-vadd/compare.py | 40 + .../backend/mixed-external-vadd/golden.py | 43 + .../backend/mixed-external-vadd/kernel.pto | 93 ++ .../backend/mixed-external-vadd/launch.cpp | 34 +- .../backend/mixed-external-vadd/main.cpp | 118 ++ .../backend/mixed-external-vadd/ptoas.flags | 1 + test/vpto/scripts/run_host_vpto_validation.sh | 9 +- tools/ptoas/CMakeLists.txt | 3 +- tools/ptoas/ObjectEmission.cpp | 1044 +++++++++++++++++ tools/ptoas/ObjectEmission.h | 146 +++ tools/ptoas/VPTOFatobjEmission.cpp | 596 ---------- tools/ptoas/VPTOHostStubEmission.cpp | 64 +- tools/ptoas/VPTOHostStubEmission.h | 4 + tools/ptoas/driver.cpp | 723 ++++++++++++ tools/ptoas/ptoas.cpp | 302 ++--- tools/ptoas/ptoas.h | 70 ++ 40 files changed, 3252 insertions(+), 896 deletions(-) create mode 100644 docs/designs/module-backend-attr-and-mixed-fatobj.md create mode 100644 test/lit/vpto/backend_attr_invalid.pto create mode 100644 test/lit/vpto/backend_attr_selects_vpto.pto create mode 100644 test/lit/vpto/backend_child_attr_inherits_outer_arch.pto create mode 100644 test/lit/vpto/backend_child_attr_selects_vpto.pto create mode 100644 test/lit/vpto/backend_cli_override.pto create mode 100644 test/lit/vpto/backend_cli_override_to_vpto.pto create mode 100644 test/lit/vpto/backend_mixed_child_attrs_rejected.pto create mode 100644 test/lit/vpto/backend_mixed_child_missing_backend_defaults.pto create mode 100644 test/lit/vpto/backend_mixed_external_import_resolves.pto create mode 100644 test/lit/vpto/backend_mixed_requires_output_file.pto create mode 100644 test/lit/vpto/backend_public_abi_export_suffix.pto create mode 100644 test/lit/vpto/backend_vpto_child_missing_kernel_kind_invalid.pto create mode 100644 test/lit/vpto/kernel_attr_vpto_llvm.pto create mode 100644 test/vpto/cases/micro-op/backend/mixed-external-vadd/compare.py create mode 100644 test/vpto/cases/micro-op/backend/mixed-external-vadd/golden.py create mode 100644 test/vpto/cases/micro-op/backend/mixed-external-vadd/kernel.pto rename tools/ptoas/VPTOFatobjEmission.h => test/vpto/cases/micro-op/backend/mixed-external-vadd/launch.cpp (53%) create mode 100644 test/vpto/cases/micro-op/backend/mixed-external-vadd/main.cpp create mode 100644 test/vpto/cases/micro-op/backend/mixed-external-vadd/ptoas.flags create mode 100644 tools/ptoas/ObjectEmission.cpp create mode 100644 tools/ptoas/ObjectEmission.h delete mode 100644 tools/ptoas/VPTOFatobjEmission.cpp create mode 100644 tools/ptoas/driver.cpp create mode 100644 tools/ptoas/ptoas.h 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..a4e16f3fb6 --- /dev/null +++ b/docs/designs/module-backend-attr-and-mixed-fatobj.md @@ -0,0 +1,442 @@ +# 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, +target-arch resolution, textual PTO parsing, PTOBC decoding, backend job +dispatch, object emission, and final output handling. + +The PTO compiler pipelines remain in `tools/ptoas/ptoas.cpp`. The driver calls +those pipelines through `compilePTOASModule`; it does not duplicate or own the +lowering pipeline itself. + +The design also replaces backend-specific object/fatobj emitters with one +`ObjectEmission` module. `ObjectEmission` provides shared CCE/Bisheng-facing +helpers for compiling C++ or VPTO LLVM artifacts and packaging or linking +fatobj objects. + +## 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 a single backend for the input. +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`. + +For a container with child modules, each child has an effective backend. A +missing child `pto.backend` uses the default backend. If all effective child +backends are the same, PTOAS uses a single backend job. If child modules have +more than one effective backend and `--pto-backend` is absent, PTOAS enters +mixed fatobj mode. + +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 VPTO module: + +```mlir +module attributes {pto.target_arch = "a5", pto.backend = "vpto"} { + func.func @kernel(...) attributes {pto.kernel} { + pto.section.vector { + ... + } + return + } +} +``` + +If `--pto-backend` is absent, this is equivalent to: + +```bash +ptoas --pto-backend=vpto input.pto -o kernel.o +``` + +`pto.kernel` is the preferred spelling for launched device kernels. The legacy +`pto.aicore` spelling is still accepted for compatibility. + +For `pto.backend = "emitc"`, PTOAS uses the existing EmitC lowering to produce +CCE C++ source. In the current single-backend EmitC path, the driver writes that +text output to `-o`; it does not internally compile the single EmitC result into +a fatobj. + +For `pto.backend = "vpto"`, normal object mode produces a host-linkable fatobj +at `-o`. VPTO object output requires an explicit file path. + +### 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 private @vpto_post(%arg0: i64) + + func.func @emitc_entry() attributes {pto.kernel} { + %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 + } + } +} +``` + +The outer module carries shared attributes such as `pto.target_arch`. Each child +module carries its own backend selection. Child modules may also carry +backend-specific attributes such as `pto.kernel_kind`. When the driver detaches +a child module into a backend job, it propagates shared outer attributes to that +child unless the child already defines them. + +Mixed-backend mode produces a final fatobj. It requires an explicit `-o` file +path and rejects debug IR output modes because those modes do not produce the +child fatobjs needed by the mixed linker. + +If child modules use more than one backend and `--pto-backend` is present, the +command line wins and the input is treated as a forced single-backend +compilation. + +## Export and Symbol Contract + +For non-kernel 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 outside the current child. A +function with a body defines the symbol in the current child module; it is +exported only when it is public. + +For `pto.kernel` functions, symbol visibility is not the import/export +contract. A `pto.kernel` function is a launched device entry and follows the +backend's kernel ABI. Users write kernel function names without `_mix_aiv` or +`_mix_aic`; VPTO lowering derives those suffixes from the normalized +`pto.kernel_kind`. + +Users write source-level non-kernel function names without `.vector` or `.cube` +ABI suffixes. The object emission path derives public ABI names for VPTO LLVM +modules from the target unit: + +| `pto.kernel_kind` | Source symbol | Generated ABI export symbol | +|-------------------|---------------|-----------------------------| +| `#pto.kernel_kind` | `@foo` | `@foo.vector` | +| `#pto.kernel_kind` | `@foo` | `@foo.cube` | + +This ABI naming is applied to compiled artifacts. The driver does not build a +separate export/import table and does not pre-resolve every cross-backend call +before compilation. + +### Mixed-Backend External Call Case + +Cross-backend calls are represented as normal external symbol references inside +the caller module. The callee is written as a source-level symbol and provided +by another backend child module or by a later link input. + +```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 @emitc_entry( + %src: !pto.ptr, + %dst: !pto.ptr, + %n: index) attributes {pto.kernel} { + 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 a public +definition for the same source symbol. The input does not spell +`@vpto_post.vector`; the backend object path derives the required ABI symbol. + +The backend pipelines do not inline or lower across the child-module boundary. +Each child is compiled into a fatobj, and the final mixed link resolves +cross-child device references with the CCE fatobj linker. + +The declaration is private because current `func.func` usage represents +body-less external declarations as private symbols. + +## Validation + +PTOAS should reject: + +1. Any `pto.backend` value other than `emitc` or `vpto`. +2. A normalized VPTO child module missing `pto.kernel_kind`. +3. Mixed fatobj mode without an explicit file path passed with `-o`. +4. Mixed fatobj mode with debug IR output flags. +5. Invalid VPTO section sugar, including duplicate vector/cube sections in one + function and nested vector/cube sections. +6. VPTO host stub generation where vector/cube variants of the same logical + kernel have incompatible signatures inside the module being compiled. + +External symbol problems that survive backend compilation are diagnosed by the +object-emission or fatobj-link stages. + +## Architecture + +```text +ptoas main (tools/ptoas/driver.cpp) + ├─ register PTOAS dialects, passes, and command-line options + ├─ parse command line + ├─ load input buffer + ├─ parse .pto or decode .ptobc + ├─ resolve target arch and backend mode + ├─ runPTOASJobs() + │ ├─ EmitCBackendJob ──────────────── compilePTOASModule -> C++ source + │ ├─ VPTOBackendJob ──────────────── compilePTOASModule -> fatobj + │ └─ mixed backend + │ ├─ EmitCBackendChildJob ───── compilePTOASModule -> child fatobj + │ ├─ VPTOBackendChildJob ───── compilePTOASModule -> child fatobj + │ └─ FatobjLinkJob ─────────── link child fatobjs + └─ write text output or keep final fatobj at -o +``` + +Pass pipelines stop at compiler artifacts: + +| Backend path | Pipeline output | +|--------------|-----------------| +| EmitC | CCE C++ source | +| VPTO | VPTO LLVM modules plus optional host stub source | + +`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` API to +call after the relevant pipeline has returned its artifact. + +## PTOAS Driver + +`ptoas` enters the driver layer immediately. Command-line parsing, input +loading, MLIR context setup, textual `.pto` parsing, and `.ptobc` decoding are +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. + +`tools/ptoas/driver.cpp` provides `main`. `tools/ptoas/ptoas.cpp` keeps the PTO +compiler options, dialect/pass registration helpers, and `compilePTOASModule`. + +### Driver Responsibilities + +The driver owns the control flow: + +```text +load input + -> parse/decode module + -> resolve backend mode + -> run backend job(s) + -> write text output or final fatobj +``` + +1. Input setup. + - Parse PTOAS command-line options. + - Track whether `--pto-backend` and `--pto-arch` appeared on the command + line. + - Load `.pto` text or `.ptobc` bytes. + - Decode `.ptobc` inputs. + - Parse textual `.pto` inputs with the effective parser target arch. + - Set or preserve `pto.target_arch`. + +2. Backend resolution. + - Parse and validate `pto.backend`. + - Decide single-backend EmitC, single-backend VPTO, or mixed-backend fatobj + mode. + - Detach backend child modules for mixed mode. + - Preserve shared outer attributes such as `pto.target_arch`. + +3. Backend jobs. + - `EmitCBackendJob` compiles a module to CCE C++ source. + - `VPTOBackendJob` compiles a module to VPTO LLVM artifacts and emits a + fatobj. + - `EmitCBackendChildJob` compiles an EmitC child to a temporary fatobj. + - `VPTOBackendChildJob` compiles a VPTO child to a temporary fatobj. + - `FatobjLinkJob` links child fatobjs into the final mixed fatobj. + +4. Output handling. + - In text-output mode, write the compiler text result to `-o`. + - In object-output mode, require an explicit file path passed with `-o`. + - In mixed mode, reject debug IR output flags before object emission starts. + +Current implementation status: + +- The driver parses `pto.backend`, detaches backend child modules, propagates + shared outer attributes, and dispatches single or mixed backend jobs. +- `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 and are resolved by the object/fatobj link flow, not + by a separate driver symbol-resolution table. +- Each VPTO child job emits its own host stub source when needed; the current + mixed path does not merge multiple VPTO children into one shared stub context. + +### Driver Data Model + +The driver keeps only the state needed to run backend jobs: + +```text +PTOASContext + MLIRContext + output path + target arch + original argc / argv + ObjectEmissionToolchain + TempFileRegistry + +PTOASCompileResult + kind: text | VPTO object | mixed object + text output + VPTO cube LLVM module + VPTO vector LLVM module + VPTO host stub source +``` + +Mixed mode does not require a separate persistent plan object. Child modules +are detached into backend jobs, and each job owns the actions required to +produce its artifact. + +## `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 composed helpers and stage-level operations. + +`ObjectEmission` does not decide backend selection and does not run PTO or VPTO +MLIR lowering pipelines. The driver produces C++ source or VPTO LLVM modules, +then requests object/fatobj operations from this component. + +### High-Level Emit Interfaces + +```text +emitFatobjCCE(cppSource) -> fatobj +emitFatobjLLVM(cubeLLVM, vectorLLVM, stubSource, moduleId) -> fatobj +emitFatobjLLVMWithRuntime(cubeLLVM, vectorLLVM, stubSource) -> fatobj +``` + +The C++ source path compiles EmitC-generated CCE source into a fatobj: + +```text +CCE C++ source + -> Bisheng CCE compilation + -> fatobj +``` + +The VPTO path compiles VPTO LLVM modules for the matching device target: + +```text +VPTO vector LLVM -> Bisheng IR compilation -> vector device object +VPTO cube LLVM -> Bisheng IR compilation -> cube device object +device object(s) + stub.cpp -> fatobj +``` + +### Fine-Grained Stage API + +`ObjectEmission` exposes stage-level APIs so the driver and tests can run +individual pieces without going through a monolithic helper: + +```text +writeCppSource(cppSource, path) +writeLLVMModule(llvmModule, path) +writeHostStubSource(stubSource, path) + +compileCppToDeviceObject(cppPath, outObjPath, target) +compileLLVMToDeviceObject(llPath, outObjPath, target) + +mergeDeviceObjects(objectPaths, outObjPath) +compileStubToFatobj(stubPath, deviceObjPath, outputPath, moduleId) +linkFatobjs(fatobjPaths, outputPath) +``` + +`compileStubToFatobj` is the stage that consumes `stub.cpp` and the compiled +device object and produces a fatobj. The current flow does not expose host stub +compilation as a separate output object. + +### 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 fatobj output. +3. Compile CCE C++ source into fatobj artifacts. +4. Compile VPTO LLVM modules into vector or cube device objects. +5. Merge VPTO device objects when both cube and vector modules are present. +6. Compile the host stub together with the device object into a fatobj. +7. Link multiple child fatobjs into the final mixed-backend fatobj. +8. Keep diagnostics separated by stage and artifact kind. + +## Implementation Status + +Implemented behavior covered by tests: + +1. `pto.backend` attr fallback when `--pto-backend` is absent. +2. CLI backend override when `--pto-backend` is present. +3. Mixed backend mode selection from child module backends. +4. Missing child `pto.backend` defaulting to the default backend. +5. VPTO child `pto.kernel_kind` validation. +6. Mixed mode requiring an explicit output file. +7. Mixed mode rejecting debug IR output flags. +8. VPTO public non-kernel ABI suffix generation from suffix-free source + symbols in object emission. +9. VPTO kernel `_mix_aiv` / `_mix_aic` suffix generation from suffix-free + `pto.kernel` source symbols. +10. Cross-backend external calls represented as normal `func.func` external + declarations and resolved by the fatobj link flow. diff --git a/include/PTO/IR/PTO.h b/include/PTO/IR/PTO.h index 981d353d1e..58642cef27 100644 --- a/include/PTO/IR/PTO.h +++ b/include/PTO/IR/PTO.h @@ -180,12 +180,20 @@ class ScopedPTOParserTargetArch { /// Function attribute that marks an explicit PTO kernel entry. inline constexpr llvm::StringLiteral kPTOEntryAttrName = "pto.entry"; inline constexpr llvm::StringLiteral kLegacyHACCEntryAttrName = "hacc.entry"; +inline constexpr llvm::StringLiteral kPTOKernelAttrName = "pto.kernel"; +inline constexpr llvm::StringLiteral kLegacyPTOAICoreAttrName = "pto.aicore"; inline constexpr llvm::StringLiteral kPTOSimtEntryAttrName = "pto.simt_entry"; inline constexpr llvm::StringLiteral kPTOSimtMaxThreadsAttrName = "pto.simt_max_threads"; inline constexpr llvm::StringLiteral kPTOSimtMaxRegistersAttrName = "pto.simt_max_regs"; +/// Return true if the operation carries a PTO kernel marker. +bool hasPTOKernelAttr(Operation *op); + +/// Return true if the function is a PTO kernel definition. +bool isPTOKernelFunction(func::FuncOp func); + /// Return true if the function carries an explicit entry marker. bool hasExplicitPTOEntryAttr(func::FuncOp func); diff --git a/include/PTO/Transforms/VPTOLLVMEmitter.h b/include/PTO/Transforms/VPTOLLVMEmitter.h index b831ac33a8..d3c1b8decc 100644 --- a/include/PTO/Transforms/VPTOLLVMEmitter.h +++ b/include/PTO/Transforms/VPTOLLVMEmitter.h @@ -37,6 +37,11 @@ struct VPTOEmissionOptions { }; struct EmittedLLVMModule { + void reset() { + module.reset(); + context.reset(); + } + std::unique_ptr context; std::unique_ptr module; }; diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 6952a7e288..887349f7eb 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -2125,6 +2125,15 @@ bool mlir::pto::isScalarPtrOrMemRef(Type type) { return false; } +bool mlir::pto::hasPTOKernelAttr(Operation *op) { + return op && (op->hasAttr(kPTOKernelAttrName) || + op->hasAttr(kLegacyPTOAICoreAttrName)); +} + +bool mlir::pto::isPTOKernelFunction(func::FuncOp func) { + return func && !func.isDeclaration() && hasPTOKernelAttr(func.getOperation()); +} + bool mlir::pto::hasExplicitPTOEntryAttr(func::FuncOp func) { return func && (func->hasAttrOfType(kPTOEntryAttrName) || func->hasAttrOfType(kLegacyHACCEntryAttrName)); diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index c664712446..e4d2a8c900 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -2740,14 +2740,15 @@ struct FuncToEmitC : public OpConversionPattern { } if (op.isDeclaration()) { - emitcFunc.setSpecifiersAttr(rewriter.getStrArrayAttr({"extern"})); + emitcFunc.setSpecifiersAttr( + rewriter.getStrArrayAttr({"extern \"C\"", "AICORE"})); 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/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index c62d3340e6..f472ae080e 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -1174,9 +1174,6 @@ struct PTOViewToMemrefPass MLIRContext *ctx = &getContext(); for (auto func : mod.getOps()) { - if (func.isExternal()) continue; - - Block &entry = func.front(); auto fnTy = func.getFunctionType(); // ------------------------------------------------------------------ @@ -1188,6 +1185,11 @@ struct PTOViewToMemrefPass SmallVector newResults; for (Type t : fnTy.getResults()) newResults.push_back(convertPTOTypeToMemRef(t)); + func.setFunctionType(FunctionType::get(ctx, newInputs, newResults)); + if (func.isExternal()) continue; + + Block &entry = func.front(); + // Update entry block arguments for (unsigned i = 0; i < entry.getNumArguments(); ++i) { if (entry.getArgument(i).getType() != newInputs[i]) { @@ -1195,9 +1197,6 @@ struct PTOViewToMemrefPass } } - // Update function type - func.setFunctionType(FunctionType::get(ctx, newInputs, newResults)); - // ------------------------------------------------------------------ // Stage 0.25: Insert pto.bind_tile for function args that were tile_buf // ------------------------------------------------------------------ diff --git a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp index 8291521580..bad9e5e695 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp @@ -59,16 +59,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); @@ -9668,7 +9661,7 @@ static LogicalResult renameKernelFunctionsForKernelKind(ModuleOp module, } for (func::FuncOp funcOp : module.getOps()) { - if (!hasVPTOKernelAttr(funcOp)) + if (!pto::hasPTOKernelAttr(funcOp.getOperation())) 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..8367461871 --- /dev/null +++ b/test/lit/vpto/backend_mixed_child_attrs_rejected.pto @@ -0,0 +1,28 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a5 --emit-vpto %s -o - 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.backend = "emitc"} { + func.func @emitc_child() attributes {pto.aicore} { + return + } + } + + module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind + } { + func.func @vpto_child() { + return + } + } +} + +// CHECK: mixed pto.backend fatobj mode does not support debug IR output flags diff --git a/test/lit/vpto/backend_mixed_child_missing_backend_defaults.pto b/test/lit/vpto/backend_mixed_child_missing_backend_defaults.pto new file mode 100644 index 0000000000..2b3f58cd4e --- /dev/null +++ b/test/lit/vpto/backend_mixed_child_missing_backend_defaults.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 { + func.func @missing_backend_child_defaults_to_emitc() 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_external_import_resolves.pto b/test/lit/vpto/backend_mixed_external_import_resolves.pto new file mode 100644 index 0000000000..d26b6dbb18 --- /dev/null +++ b/test/lit/vpto/backend_mixed_external_import_resolves.pto @@ -0,0 +1,35 @@ +// 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: ASCEND_HOME_PATH is required 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_public_abi_export_suffix.pto b/test/lit/vpto/backend_public_abi_export_suffix.pto new file mode 100644 index 0000000000..f13e0e4fa5 --- /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 +// CHECK: call @local_helper +// CHECK: func.func private @local_helper +// CHECK: func.func @device_kernel +// CHECK: call @vpto_post +// CHECK-NOT: @vpto_post.vector diff --git a/test/lit/vpto/backend_vpto_child_missing_kernel_kind_invalid.pto b/test/lit/vpto/backend_vpto_child_missing_kernel_kind_invalid.pto new file mode 100644 index 0000000000..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/kernel_attr_vpto_llvm.pto b/test/lit/vpto/kernel_attr_vpto_llvm.pto new file mode 100644 index 0000000000..a6a0b36c36 --- /dev/null +++ b/test/lit/vpto/kernel_attr_vpto_llvm.pto @@ -0,0 +1,17 @@ +// 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: ( mkdir -p %T && ptoas --pto-arch=a5 --pto-backend=vpto %s -o %t --mlir-print-ir-after=convert-func-to-llvm 2>&1 || true ) | FileCheck %s + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @kernel_attr_entry() attributes {pto.kernel} { + return + } +} + +// CHECK-LABEL: llvm.func @kernel_attr_entry_mix_aiv diff --git a/test/lit/vpto/legacy_aicore_kernel_attr.pto b/test/lit/vpto/legacy_aicore_kernel_attr.pto index e699f5c1c1..25a7c6e386 100644 --- a/test/lit/vpto/legacy_aicore_kernel_attr.pto +++ b/test/lit/vpto/legacy_aicore_kernel_attr.pto @@ -24,10 +24,10 @@ module attributes {pto.target_arch = "a5"} { } } -// SPLIT: module attributes {pto.kernel_kind = #pto.kernel_kind +// SPLIT: module attributes {{.*}}pto.kernel_kind = #pto.kernel_kind // SPLIT: func.func @legacy_attr_kernel // SPLIT-SAME: attributes {pto.aicore} -// SPLIT: module attributes {pto.kernel_kind = #pto.kernel_kind +// SPLIT: module attributes {{.*}}pto.kernel_kind = #pto.kernel_kind // SPLIT: func.func @legacy_attr_kernel // SPLIT-SAME: attributes {pto.aicore} diff --git a/test/lit/vpto/section_sugar_mixed.pto b/test/lit/vpto/section_sugar_mixed.pto index f855ef6196..3ff8fcf892 100644 --- a/test/lit/vpto/section_sugar_mixed.pto +++ b/test/lit/vpto/section_sugar_mixed.pto @@ -23,7 +23,8 @@ module attributes {pto.target_arch = "a5"} { } } -// CHECK: module attributes {pto.target_arch = "a5"} { +// CHECK: module attributes +// CHECK-SAME: pto.target_arch = "a5" // CHECK: module attributes { // CHECK-SAME: pto.kernel_kind = #pto.kernel_kind // CHECK: func.func @section_sugar_kernel diff --git a/test/lit/vpto/section_sugar_multi_func.pto b/test/lit/vpto/section_sugar_multi_func.pto index a0f87ff71f..5d790309af 100644 --- a/test/lit/vpto/section_sugar_multi_func.pto +++ b/test/lit/vpto/section_sugar_multi_func.pto @@ -14,7 +14,8 @@ module attributes {pto.target_arch = "a5"} { } } -// CHECK: module attributes {pto.target_arch = "a5"} { +// CHECK: module attributes +// CHECK-SAME: pto.target_arch = "a5" // CHECK: module attributes { // CHECK-SAME: pto.kernel_kind = #pto.kernel_kind // CHECK: func.func @vec_only diff --git a/test/vpto/cases/micro-op/backend/mixed-external-vadd/compare.py b/test/vpto/cases/micro-op/backend/mixed-external-vadd/compare.py new file mode 100644 index 0000000000..792d7644d5 --- /dev/null +++ b/test/vpto/cases/micro-op/backend/mixed-external-vadd/compare.py @@ -0,0 +1,40 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +import os +import sys + +import numpy as np + + +def compare_bin(golden_path, output_path, dtype, eps): + if not os.path.exists(golden_path) or not os.path.exists(output_path): + return False + golden = np.fromfile(golden_path, dtype=dtype) + output = np.fromfile(output_path, dtype=dtype) + return ( + golden.shape == output.shape + and np.allclose(golden, output, atol=eps, rtol=eps, equal_nan=True) + ) + + +def main(): + strict = os.getenv("COMPARE_STRICT", "1") != "0" + ok = compare_bin("golden_v3.bin", "v3.bin", np.float32, 1e-4) + if not ok: + if strict: + print("[ERROR] compare failed") + sys.exit(2) + print("[WARN] compare failed (non-gating)") + return + print("[INFO] compare passed") + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/micro-op/backend/mixed-external-vadd/golden.py b/test/vpto/cases/micro-op/backend/mixed-external-vadd/golden.py new file mode 100644 index 0000000000..af261b5830 --- /dev/null +++ b/test/vpto/cases/micro-op/backend/mixed-external-vadd/golden.py @@ -0,0 +1,43 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +import argparse +from pathlib import Path + +import numpy as np + + +ELEM_COUNT = 1024 +SEED = 29 + + +def generate(output_dir: Path, seed: int) -> None: + rng = np.random.default_rng(seed) + lhs = rng.uniform(-8.0, 8.0, size=(ELEM_COUNT,)).astype(np.float32) + rhs = rng.uniform(-8.0, 8.0, size=(ELEM_COUNT,)).astype(np.float32) + out = np.zeros((ELEM_COUNT,), dtype=np.float32) + golden = (lhs + rhs).astype(np.float32, copy=False) + + output_dir.mkdir(parents=True, exist_ok=True) + lhs.tofile(output_dir / "v1.bin") + rhs.tofile(output_dir / "v2.bin") + out.tofile(output_dir / "v3.bin") + golden.tofile(output_dir / "golden_v3.bin") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir", type=Path, default=Path(".")) + parser.add_argument("--seed", type=int, default=SEED) + args = parser.parse_args() + generate(args.output_dir, args.seed) + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/micro-op/backend/mixed-external-vadd/kernel.pto b/test/vpto/cases/micro-op/backend/mixed-external-vadd/kernel.pto new file mode 100644 index 0000000000..197ace46c1 --- /dev/null +++ b/test/vpto/cases/micro-op/backend/mixed-external-vadd/kernel.pto @@ -0,0 +1,93 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.backend = "emitc"} { + func.func private @external_vadd(%lhs: !pto.ptr, + %rhs: !pto.ptr, + %out: !pto.ptr) + + func.func @mixed_external_vadd_kernel(%lhs: !pto.ptr, + %rhs: !pto.ptr, + %out: !pto.ptr) + attributes {pto.aicore} { + pto.section.vector { + func.call @external_vadd(%lhs, %rhs, %out) + : (!pto.ptr, !pto.ptr, !pto.ptr) -> () + } + return + } + } + + module attributes { + pto.backend = "vpto", + pto.kernel_kind = #pto.kernel_kind + } { + func.func public @external_vadd(%lhs_gm: !pto.ptr, + %rhs_gm: !pto.ptr, + %out_gm: !pto.ptr) { + %c0 = arith.constant 0 : index + %c64 = arith.constant 64 : index + %c1024 = arith.constant 1024 : index + %c0_i64 = arith.constant 0 : i64 + %c32_i64 = arith.constant 32 : i64 + %c128_i64 = arith.constant 128 : i64 + %c4096_i64 = arith.constant 4096 : i64 + %c8192_i64 = arith.constant 8192 : i64 + %c1024_i32 = arith.constant 1024 : i32 + + %ub_lhs = pto.castptr %c0_i64 : i64 -> !pto.ptr + %ub_rhs = pto.castptr %c4096_i64 : i64 -> !pto.ptr + %ub_out = pto.castptr %c8192_i64 : i64 -> !pto.ptr + + pto.get_buf "PIPE_MTE2", 0, 0 + pto.mte_gm_ub %lhs_gm, %ub_lhs, %c0_i64, %c128_i64 + nburst(%c32_i64, %c128_i64, %c128_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.rls_buf "PIPE_MTE2", 0, 0 + + pto.get_buf "PIPE_MTE2", 1, 0 + pto.mte_gm_ub %rhs_gm, %ub_rhs, %c0_i64, %c128_i64 + nburst(%c32_i64, %c128_i64, %c128_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.rls_buf "PIPE_MTE2", 1, 0 + + pto.get_buf "PIPE_V", 0, 0 + pto.get_buf "PIPE_V", 1, 0 + pto.get_buf "PIPE_V", 2, 0 + pto.vecscope { + %_:1 = scf.for %offset = %c0 to %c1024 step %c64 + iter_args(%remaining = %c1024_i32) -> (i32) { + %mask, %next_remaining = pto.plt_b32 %remaining + : i32 -> !pto.mask, i32 + %lhs = pto.vlds %ub_lhs[%offset] + : !pto.ptr -> !pto.vreg<64xf32> + %rhs = pto.vlds %ub_rhs[%offset] + : !pto.ptr -> !pto.vreg<64xf32> + %sum = pto.vadd %lhs, %rhs, %mask + : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask + -> !pto.vreg<64xf32> + pto.vsts %sum, %ub_out[%offset], %mask + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + scf.yield %next_remaining : i32 + } + } + pto.rls_buf "PIPE_V", 0, 0 + pto.rls_buf "PIPE_V", 1, 0 + pto.rls_buf "PIPE_V", 2, 0 + + pto.get_buf "PIPE_MTE3", 2, 0 + pto.mte_ub_gm %ub_out, %out_gm, %c128_i64 + nburst(%c32_i64, %c128_i64, %c128_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.rls_buf "PIPE_MTE3", 2, 0 + pto.barrier #pto.pipe + return + } + } +} diff --git a/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 5763019fa9..709b29747f 100755 --- a/test/vpto/scripts/run_host_vpto_validation.sh +++ b/test/vpto/scripts/run_host_vpto_validation.sh @@ -255,6 +255,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}" @@ -263,8 +264,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/ObjectEmission.cpp b/tools/ptoas/ObjectEmission.cpp new file mode 100644 index 0000000000..c6976687b1 --- /dev/null +++ b/tools/ptoas/ObjectEmission.cpp @@ -0,0 +1,1044 @@ +// 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 "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" +#include "llvm/Support/MemoryBuffer.h" +#include "llvm/Support/Path.h" +#include "llvm/Support/Process.h" +#include "llvm/Support/Program.h" +#include "llvm/Support/ToolOutputFile.h" +#include "llvm/TargetParser/Host.h" +#include "llvm/Support/raw_ostream.h" + +#include +#include +#include +#include + +namespace { + +using llvm::StringRef; + +static bool runCommandWithStderr(llvm::StringRef program, + llvm::ArrayRef ownedArgs, + llvm::StringRef stderrPath, + llvm::raw_ostream &diagOS, + llvm::StringRef what, + std::optional stdinPath = + std::nullopt); + +static bool writeTextFile(StringRef path, StringRef content, + llvm::raw_ostream &diagOS) { + std::error_code ec; + llvm::raw_fd_ostream os(path, ec, llvm::sys::fs::OF_Text); + if (ec) { + diagOS << "Error: failed to open " << path << " for write: " + << ec.message() << "\n"; + return false; + } + os << content; + os.flush(); + return true; +} + +static void stripUnsupportedBishengAttrs(llvm::Module &module) { + for (llvm::Function &function : module) { + // LLVM 19 prints memory effect attributes in textual form like + // `memory(none)`. beta.1 Bisheng cannot parse that syntax, so remove only + // the unsupported memory-effect attribute before serializing the module. + function.setAttributes( + function.getAttributes().removeFnAttribute(module.getContext(), + llvm::Attribute::Memory)); + } +} + +static bool writeLLVMModuleFile(llvm::Module &module, StringRef path, + llvm::raw_ostream &diagOS) { + std::error_code ec; + llvm::raw_fd_ostream os(path, ec, llvm::sys::fs::OF_Text); + if (ec) { + diagOS << "Error: failed to open " << path << " for write: " + << ec.message() << "\n"; + return false; + } + stripUnsupportedBishengAttrs(module); + module.print(os, nullptr); + os.flush(); + return true; +} + +static std::string sanitizeModuleId(llvm::StringRef raw) { + std::string out; + out.reserve(raw.size()); + for (char c : raw) { + if (std::isalnum(static_cast(c)) || c == '_') + out.push_back(c); + else + out.push_back('_'); + } + if (out.empty()) + out = "ptoas_fatobj"; + return out; +} + +static std::optional getAscendHomePath() { + const char *env = std::getenv("ASCEND_HOME_PATH"); + if (!env || !*env) + return std::nullopt; + 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); + return std::string(joined.str()); +} + +static std::optional locateProgram(llvm::StringRef envPath, + llvm::StringRef fallbackName) { + if (!envPath.empty() && llvm::sys::fs::exists(envPath)) + return envPath.str(); + if (auto found = llvm::sys::findProgramByName(fallbackName)) + return *found; + 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, + llvm::StringRef outObjPath, + llvm::StringRef targetCPU, + llvm::StringRef bishengPath, + llvm::StringRef stderrPath, + llvm::raw_ostream &diagOS); +static bool compileHostStubToObject(llvm::StringRef stubPath, + llvm::StringRef outObjPath, + llvm::StringRef moduleId, + llvm::StringRef targetCPU, + const VPTOFatobjToolchain &toolchain, + llvm::StringRef deviceObjPath, + llvm::StringRef stderrPath, + llvm::raw_ostream &diagOS); +static bool mergeDeviceObjects(llvm::ArrayRef deviceObjPaths, + llvm::StringRef outObjPath, + llvm::StringRef ldLldPath, + 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(); + if (!ascendHome) { + diagOS << "Error: ASCEND_HOME_PATH is required for VPTO fatobj emission.\n"; + return std::nullopt; + } + + VPTOFatobjToolchain toolchain(*ascendHome, diagOS); + if (!toolchain.validate(diagOS)) + return std::nullopt; + return toolchain; + } + + const std::string &ascendHome() const { return ascendHomePath; } + const std::string &bisheng() const { return bishengPath; } + const std::string &bishengCc1() const { return bishengCc1Path; } + const std::string &cceLd() const { return cceLdPath; } + const std::string &ldLld() const { return ldLldPath; } + const std::string &resourceDir() const { return resourceDirPath; } + const std::string &resourceIncludeDir() const { + return resourceIncludeDirPath; + } + const std::string &cceStubDir() const { return cceStubDirPath; } + const std::string &bishengCompilerBinDir() const { + 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, + llvm::raw_ostream &diagOS) + : ascendHomePath(ascendHome.str()), + bishengPath(joinPath(ascendHomePath, "bin/bisheng")), + bishengCc1Path( + joinPath(ascendHomePath, "tools/bisheng_compiler/bin/bisheng")), + cceLdPath(joinPath(ascendHomePath, "bin/cce-ld")), + ldLldPath( + locateProgram(joinPath(ascendHomePath, "bin/ld.lld"), "ld.lld") + .value_or(std::string())), + resourceDirPath(joinPath( + ascendHomePath, "tools/bisheng_compiler/lib/clang/15.0.5")), + resourceIncludeDirPath(joinPath(resourceDirPath, "include")), + cceStubDirPath(joinPath(resourceIncludeDirPath, "cce_stub")), + bishengCompilerBinDirPath( + joinPath(ascendHomePath, "tools/bisheng_compiler/bin")) { + cppIncludeDirs = discoverCppIncludeDirs(ascendHomePath, diagOS, + ptoIsaPath); + } + + bool validate(llvm::raw_ostream &diagOS) const { + if (!llvm::sys::fs::exists(bishengPath)) { + diagOS << "Error: unable to locate bisheng: " << bishengPath << "\n"; + return false; + } + if (!llvm::sys::fs::exists(bishengCc1Path)) { + diagOS << "Error: unable to locate bisheng cc1 frontend: " + << bishengCc1Path << "\n"; + return false; + } + if (!llvm::sys::fs::exists(cceLdPath)) { + diagOS << "Error: unable to locate cce-ld: " << cceLdPath << "\n"; + return false; + } + if (ldLldPath.empty() || !llvm::sys::fs::exists(ldLldPath)) { + diagOS << "Error: unable to locate ld.lld.\n"; + return false; + } + return true; + } + + 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; + llvm::SmallVector cppIncludeDirs; +}; + +class VPTOFatobjArtifacts { +public: + explicit VPTOFatobjArtifacts(mlir::pto::TempFileRegistry &tempFiles) + : tempFiles(tempFiles) {} + + bool emitStubSource(StringRef stubSource, llvm::raw_ostream &diagOS) { + if (failed(tempFiles.create("ptoas-host-stub", ".cpp", stubPath, diagOS))) + return false; + if (!writeTextFile(stubPath, stubSource, diagOS)) + return false; + return true; + } + + bool initCommandLogs(llvm::raw_ostream &diagOS) { + if (failed(tempFiles.create("ptoas-stderr", ".log", stderrPath, diagOS))) + return false; + return true; + } + + bool emitCubeObject(llvm::Module *module, + const VPTOFatobjToolchain &toolchain, + llvm::raw_ostream &diagOS) { + if (!module) + return true; + if (failed(tempFiles.create("ptoas-device", ".ll", cubeLLPath, diagOS))) + return false; + if (failed(tempFiles.create("ptoas-device", ".o", cubeObjPath, diagOS))) + return false; + return succeeded(mlir::pto::emitVPTOCubeDeviceObject( + *module, cubeLLPath, cubeObjPath, toolchain.toPublicToolchain(), + stderrPath, diagOS)); + } + + bool emitVectorObject(llvm::Module *module, + const VPTOFatobjToolchain &toolchain, + llvm::raw_ostream &diagOS) { + if (!module) + return true; + if (failed(tempFiles.create("ptoas-device", ".ll", vectorLLPath, diagOS))) + return false; + if (failed(tempFiles.create("ptoas-device", ".o", vectorObjPath, diagOS))) + return false; + return succeeded(mlir::pto::emitVPTOVectorDeviceObject( + *module, vectorLLPath, vectorObjPath, toolchain.toPublicToolchain(), + stderrPath, diagOS)); + } + + bool mergeDeviceObjects(const VPTOFatobjToolchain &toolchain, + llvm::raw_ostream &diagOS) { + llvm::SmallVector deviceObjPaths; + if (!cubeObjPath.empty()) + deviceObjPaths.push_back(cubeObjPath); + if (!vectorObjPath.empty()) + deviceObjPaths.push_back(vectorObjPath); + if (deviceObjPaths.empty()) { + diagOS << "Error: VPTO fatobj emission requires at least one device module.\n"; + return false; + } + if (failed(tempFiles.create("ptoas-device-merged", ".o", + mergedDeviceObjPath, diagOS))) + return false; + return ::mergeDeviceObjects(deviceObjPaths, mergedDeviceObjPath, + toolchain.ldLld(), stderrPath, diagOS); + } + + bool compileHostStub(const VPTOFatobjToolchain &toolchain, + llvm::StringRef moduleId, + llvm::StringRef targetCPU, + llvm::raw_ostream &diagOS) { + if (failed(tempFiles.create("ptoas-host-stub", ".o", hostStubObjPath, + diagOS))) + return false; + return compileHostStubToObject(stubPath, hostStubObjPath, moduleId, + targetCPU, toolchain, mergedDeviceObjPath, + stderrPath, diagOS); + } + + bool compileHostStubToFatobj(const VPTOFatobjToolchain &toolchain, + llvm::StringRef moduleId, + llvm::StringRef targetCPU, + llvm::StringRef outputPath, + llvm::raw_ostream &diagOS) { + return compileHostStubToObject(stubPath, outputPath, moduleId, targetCPU, + toolchain, mergedDeviceObjPath, stderrPath, + diagOS); + } + + bool repackFatObj(const VPTOFatobjToolchain &toolchain, + llvm::StringRef moduleId, llvm::StringRef targetCPU, + llvm::StringRef outPath, llvm::raw_ostream &diagOS) { + llvm::SmallVector args = { + toolchain.cceLd(), + toolchain.ldLld(), + "-x", + "-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(), + "-cce-stub-dir", + toolchain.cceStubDir(), + "-cce-install-dir", + toolchain.bishengCompilerBinDir(), + "-cce-inputs-number", + "1", + hostStubObjPath, + }; + return runCommandWithStderr(toolchain.cceLd(), args, stderrPath, diagOS, + "fatobj repack"); + } + +private: + mlir::pto::TempFileRegistry &tempFiles; + std::string cubeLLPath; + std::string cubeObjPath; + std::string vectorLLPath; + std::string vectorObjPath; + std::string mergedDeviceObjPath; + std::string stderrPath; + std::string stubPath; + std::string hostStubObjPath; +}; + +static bool runCommandWithStderr(llvm::StringRef program, + llvm::ArrayRef ownedArgs, + llvm::StringRef stderrPath, + llvm::raw_ostream &diagOS, + llvm::StringRef what, + std::optional stdinPath) { + llvm::SmallVector args; + args.reserve(ownedArgs.size()); + for (const std::string &arg : ownedArgs) + args.push_back(arg); + llvm::SmallVector, 3> redirects = { + stdinPath, std::nullopt, stderrPath}; + + std::string execErr; + bool execFailed = false; + int rc = llvm::sys::ExecuteAndWait(program, args, std::nullopt, redirects, 0, + 0, &execErr, &execFailed); + if (!execFailed && rc == 0) + return true; + + diagOS << "Error: " << what << " failed\n"; + diagOS << "Command:"; + for (llvm::StringRef arg : args) + diagOS << " " << arg; + diagOS << "\n"; + if (!execErr.empty()) + diagOS << execErr << "\n"; + if (auto buffer = llvm::MemoryBuffer::getFile(stderrPath)) + diagOS << buffer.get()->getBuffer() << "\n"; + return false; +} + +static bool compileDeviceLLVMToObject(llvm::StringRef llPath, + llvm::StringRef outObjPath, + llvm::StringRef targetCPU, + llvm::StringRef bishengPath, + llvm::StringRef stderrPath, + llvm::raw_ostream &diagOS) { + llvm::SmallVector args = { + bishengPath.str(), + "--target=hiipu64-hisilicon-cce", + std::string("-march=") + targetCPU.str(), + std::string("--cce-aicore-arch=") + targetCPU.str(), + "--cce-aicore-only", + "-O2", + "-dc", + "-mllvm", + "-cce-dyn-kernel-stack-size=true", + "-c", + "-x", + "ir", + "-", + "-o", + outObjPath.str(), + }; + return runCommandWithStderr(bishengPath, args, stderrPath, diagOS, + "device LLVM compilation", llPath); +} + +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, + const VPTOFatobjToolchain &toolchain, + llvm::StringRef deviceObjPath, + llvm::StringRef stderrPath, + llvm::raw_ostream &diagOS) { + std::string coverageDir = "."; + std::string debugDir = "."; + std::string hostTriple = llvm::sys::getProcessTriple(); + + llvm::SmallVector args = { + toolchain.bishengCc1(), + "-cc1", + "-triple", + hostTriple, + "-target-cpu", + llvm::sys::getHostCPUName().str(), + "-fcce-aicpu-legacy-launch", + "-fcce-is-host", + "-cce-enable-mix", + "-mllvm", + "-enable-mix=true", + "-cce-launch-with-flagv2-impl", + "-fcce-aicore-arch", + targetCPU.str(), + "-fcce-fatobj-compile", + "-emit-obj", + "--mrelax-relocations", + "-disable-free", + "-clear-ast-before-backend", + "-disable-llvm-verifier", + "-discard-value-names", + "-main-file-name", + "stub.cpp", + "-mrelocation-model", + "pic", + "-pic-level", + "2", + "-fhalf-no-semantic-interposition", + "-mframe-pointer=none", + "-fmath-errno", + "-ffp-contract=on", + "-fno-rounding-math", + "-mconstructor-aliases", + "-funwind-tables=2", + "-fallow-half-arguments-and-returns", + "-mllvm", + "-treat-scalable-fixed-error-as-warning", + std::string("-fcoverage-compilation-dir=") + coverageDir, + "-resource-dir", + toolchain.resourceDir(), + "-internal-isystem", + toolchain.resourceIncludeDir(), + "-include", + "__clang_cce_runtime_wrapper.h", + "-D", + "_FORTIFY_SOURCE=2", + "-D", + "REGISTER_BASE", + "-O2", + "-Wno-macro-redefined", + "-Wno-ignored-attributes", + "-std=c++17", + "-fdeprecated-macro", + std::string("-fdebug-compilation-dir=") + debugDir, + "-ferror-limit", + "19", + "-stack-protector", + "2", + "-fno-signed-char", + "-fgnuc-version=4.2.1", + "-fcxx-exceptions", + "-fexceptions", + "-vectorize-loops", + "-vectorize-slp", + "-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", + "-fcce-include-aibinary", + deviceObjPath.str(), + "-fcce-device-module-id", + moduleId.str(), + "-faddrsig", + "-D__GCC_HAVE_DWARF2_CFI_ASM=1", + "-o", + outObjPath.str(), + "-x", + "cce", + stubPath.str(), + }; + return runCommandWithStderr(toolchain.bishengCc1(), args, stderrPath, diagOS, + "host stub compilation"); +} + +static bool mergeDeviceObjects(llvm::ArrayRef deviceObjPaths, + llvm::StringRef outObjPath, + llvm::StringRef ldLldPath, + llvm::StringRef stderrPath, + llvm::raw_ostream &diagOS) { + if (deviceObjPaths.empty()) + return false; + + llvm::SmallVector args = { + ldLldPath.str(), + "-m", + "aicorelinux", + "-Ttext", + "0", + }; + for (const std::string &path : deviceObjPaths) + args.push_back(path); + args.push_back("-o"); + args.push_back(outObjPath.str()); + args.push_back("-r"); + args.push_back("--allow-multiple-definition"); + return runCommandWithStderr(ldLldPath, args, stderrPath, diagOS, + "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::pto::TempFileRegistry::~TempFileRegistry() { cleanup(); } + +void mlir::pto::TempFileRegistry::cleanup() { + for (const std::string &path : paths) + llvm::sys::fs::remove(path); + paths.clear(); +} + +mlir::LogicalResult +mlir::pto::TempFileRegistry::create(llvm::StringRef prefix, + llvm::StringRef suffix, std::string &path, + llvm::raw_ostream &diagOS) { + llvm::SmallString<128> tempPath; + int fd = -1; + std::error_code ec = + llvm::sys::fs::createTemporaryFile(prefix, suffix, fd, tempPath); + if (ec) { + diagOS << "Error: failed to create temporary file for " << prefix << suffix + << ": " << ec.message() << "\n"; + return failure(); + } + llvm::sys::Process::SafelyCloseFileDescriptor(fd); + path = tempPath.str().str(); + paths.push_back(path); + return success(); +} + +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::emitFatobjCCE( + llvm::StringRef cppSource, llvm::StringRef outputPath, + const ObjectEmissionToolchain &toolchain, TempFileRegistry &tempFiles, + llvm::raw_ostream &diagOS) { + std::string cppPath; + std::string stderrPath; + if (failed(tempFiles.create("ptoas-emitc", ".cpp", cppPath, diagOS)) || + failed(tempFiles.create("ptoas-emitc-fatobj", ".log", stderrPath, + diagOS))) + return failure(); + return emitCppFatobj(cppSource, cppPath, outputPath, toolchain, stderrPath, + diagOS); +} + +static bool isVPTOKernelABISymbol(llvm::StringRef name) { + return name.ends_with("_mix_aiv") || name.ends_with("_mix_aic"); +} + +static bool isVPTOPublicABISymbol(llvm::StringRef name) { + return name.ends_with(".vector") || name.ends_with(".cube"); +} + +static mlir::LogicalResult renameLLVMFunction(llvm::Module &module, + llvm::StringRef sourceName, + llvm::StringRef abiName, + llvm::raw_ostream &diagOS) { + if (sourceName == abiName) + return mlir::success(); + llvm::Function *function = module.getFunction(sourceName); + if (!function) + return mlir::success(); + if (llvm::Function *existing = module.getFunction(abiName); + existing && existing != function) { + diagOS << "Error: cannot rename LLVM symbol '" << sourceName << "' to '" + << abiName << "': target symbol already exists.\n"; + return mlir::failure(); + } + function->setName(abiName); + return mlir::success(); +} + +static mlir::LogicalResult applyVPTOLLVMABINames(llvm::Module &module, + llvm::StringRef suffix, + llvm::raw_ostream &diagOS) { + for (llvm::Function &function : module) { + if (function.isDeclaration() || !function.hasExternalLinkage()) + continue; + llvm::StringRef name = function.getName(); + if (name.empty() || isVPTOKernelABISymbol(name) || + isVPTOPublicABISymbol(name)) + continue; + if (failed(renameLLVMFunction(module, name, (name + suffix).str(), diagOS))) + return mlir::failure(); + } + return mlir::success(); +} + +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(applyVPTOLLVMABINames(module, ".vector", diagOS))) + return failure(); + 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(applyVPTOLLVMABINames(module, ".cube", diagOS))) + return failure(); + if (failed(writeLLVMModule(module, llPath, diagOS))) + return failure(); + return compileLLVMToDeviceObject(llPath, outObjPath, + ObjectEmissionDeviceTarget::Cube, + toolchain, stderrPath, diagOS); +} + +mlir::LogicalResult mlir::pto::emitFatobjLLVM( + llvm::Module *cubeModule, llvm::Module *vectorModule, + llvm::StringRef stubSource, llvm::StringRef outputPath, + llvm::StringRef moduleId, const ObjectEmissionToolchain &toolchain, + TempFileRegistry &tempFiles, llvm::raw_ostream &diagOS) { + if (!cubeModule && !vectorModule) { + diagOS << "Error: VPTO fatobj emission requires at least one LLVM module.\n"; + return failure(); + } + + VPTOFatobjToolchain privateToolchain(toolchain); + VPTOFatobjArtifacts artifacts(tempFiles); + if (!artifacts.emitStubSource(stubSource, diagOS)) + return failure(); + if (!artifacts.initCommandLogs(diagOS)) + return failure(); + if (!artifacts.emitCubeObject(cubeModule, privateToolchain, diagOS)) + return failure(); + if (!artifacts.emitVectorObject(vectorModule, privateToolchain, diagOS)) + return failure(); + if (!artifacts.mergeDeviceObjects(privateToolchain, diagOS)) + return failure(); + + constexpr llvm::StringLiteral targetCPU = "dav-c310"; + if (!artifacts.compileHostStubToFatobj(privateToolchain, moduleId, targetCPU, + outputPath, diagOS)) + return failure(); + return success(); +} + +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::emitFatobjLLVMWithRuntime( + 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(); + } + + std::optional toolchain = + VPTOFatobjToolchain::create(diagOS); + if (!toolchain) + return failure(); + + TempFileRegistry tempFiles; + VPTOFatobjArtifacts artifacts(tempFiles); + if (!artifacts.emitStubSource(stubSource, diagOS)) + return failure(); + if (!artifacts.initCommandLogs(diagOS)) + return failure(); + + if (!artifacts.emitCubeObject(cubeModule, *toolchain, diagOS)) + return failure(); + if (!artifacts.emitVectorObject(vectorModule, *toolchain, diagOS)) + return failure(); + + if (!artifacts.mergeDeviceObjects(*toolchain, diagOS)) + return failure(); + + std::string moduleId = sanitizeModuleId(outputFile.getFilename()); + constexpr llvm::StringLiteral hostTargetCPU = "dav-c310"; + if (!artifacts.compileHostStub(*toolchain, moduleId, hostTargetCPU, diagOS)) + return failure(); + + if (!artifacts.repackFatObj(*toolchain, moduleId, hostTargetCPU, + outputFile.getFilename(), diagOS)) + return failure(); + outputFile.keep(); + return success(); +} diff --git a/tools/ptoas/ObjectEmission.h b/tools/ptoas/ObjectEmission.h new file mode 100644 index 0000000000..e4abc5521f --- /dev/null +++ b/tools/ptoas/ObjectEmission.h @@ -0,0 +1,146 @@ +// 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/ArrayRef.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StringRef.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; +}; + +class TempFileRegistry { +public: + ~TempFileRegistry(); + + void cleanup(); + LogicalResult create(llvm::StringRef prefix, llvm::StringRef suffix, + std::string &path, llvm::raw_ostream &diagOS); + +private: + llvm::SmallVector paths; +}; + +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 emitFatobjCCE(llvm::StringRef cppSource, + llvm::StringRef outputPath, + const ObjectEmissionToolchain &toolchain, + TempFileRegistry &tempFiles, + 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 emitFatobjLLVM( + llvm::Module *cubeModule, llvm::Module *vectorModule, + llvm::StringRef stubSource, llvm::StringRef outputPath, + llvm::StringRef moduleId, const ObjectEmissionToolchain &toolchain, + TempFileRegistry &tempFiles, 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 emitFatobjLLVMWithRuntime(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/VPTOFatobjEmission.cpp b/tools/ptoas/VPTOFatobjEmission.cpp deleted file mode 100644 index f674e3dac1..0000000000 --- a/tools/ptoas/VPTOFatobjEmission.cpp +++ /dev/null @@ -1,596 +0,0 @@ -// Copyright (c) 2026 Huawei Technologies Co., Ltd. -// This program is free software, you can redistribute it and/or modify it under the terms and conditions of -// CANN Open Software License Agreement Version 2.0 (the "License"). -// Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -// See LICENSE in the root of the software repository for the full text of the License. - -#include "VPTOFatobjEmission.h" - -#include "PTO/Transforms/VPTOLLVMEmitter.h" - -#include "llvm/ADT/SmallString.h" -#include "llvm/ADT/SmallVector.h" -#include "llvm/IR/Module.h" -#include "llvm/Support/Error.h" -#include "llvm/Support/FileSystem.h" -#include "llvm/Support/MemoryBuffer.h" -#include "llvm/Support/Path.h" -#include "llvm/Support/Process.h" -#include "llvm/Support/Program.h" -#include "llvm/Support/ToolOutputFile.h" -#include "llvm/TargetParser/Host.h" -#include "llvm/Support/raw_ostream.h" - -#include -#include -#include -#include - -namespace { - -using llvm::StringRef; - -static bool runCommandWithStderr(llvm::StringRef program, - llvm::ArrayRef ownedArgs, - llvm::StringRef stderrPath, - llvm::raw_ostream &diagOS, - llvm::StringRef what, - std::optional stdinPath = - std::nullopt); - -class TempFileRegistry { -public: - ~TempFileRegistry() { cleanup(); } - - void cleanup() { - for (const std::string &path : paths) - llvm::sys::fs::remove(path); - paths.clear(); - } - - bool create(StringRef prefix, StringRef suffix, std::string &path, - llvm::raw_ostream &diagOS) { - llvm::SmallString<128> tempPath; - int fd = -1; - std::error_code ec = llvm::sys::fs::createTemporaryFile(prefix, suffix, fd, - tempPath); - if (ec) { - diagOS << "Error: failed to create temporary file for " << prefix - << suffix << ": " << ec.message() << "\n"; - return false; - } - llvm::sys::Process::SafelyCloseFileDescriptor(fd); - path = tempPath.str().str(); - paths.push_back(path); - return true; - } - -private: - llvm::SmallVector paths; -}; - -static bool writeTextFile(StringRef path, StringRef content, - llvm::raw_ostream &diagOS) { - std::error_code ec; - llvm::raw_fd_ostream os(path, ec, llvm::sys::fs::OF_Text); - if (ec) { - diagOS << "Error: failed to open " << path << " for write: " - << ec.message() << "\n"; - return false; - } - os << content; - os.flush(); - return true; -} - -static void stripUnsupportedBishengAttrs(llvm::Module &module) { - for (llvm::Function &function : module) { - // LLVM 19 prints memory effect attributes in textual form like - // `memory(none)`. beta.1 Bisheng cannot parse that syntax, so remove only - // the unsupported memory-effect attribute before serializing the module. - function.setAttributes( - function.getAttributes().removeFnAttribute(module.getContext(), - llvm::Attribute::Memory)); - } -} - -static bool writeLLVMModuleFile(llvm::Module &module, StringRef path, - llvm::raw_ostream &diagOS) { - std::error_code ec; - llvm::raw_fd_ostream os(path, ec, llvm::sys::fs::OF_Text); - if (ec) { - diagOS << "Error: failed to open " << path << " for write: " - << ec.message() << "\n"; - return false; - } - stripUnsupportedBishengAttrs(module); - module.print(os, nullptr); - os.flush(); - return true; -} - -static std::string sanitizeModuleId(llvm::StringRef raw) { - std::string out; - out.reserve(raw.size()); - for (char c : raw) { - if (std::isalnum(static_cast(c)) || c == '_') - out.push_back(c); - else - out.push_back('_'); - } - if (out.empty()) - out = "ptoas_fatobj"; - return out; -} - -static std::optional getAscendHomePath() { - const char *env = std::getenv("ASCEND_HOME_PATH"); - 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); - return std::string(joined.str()); -} - -static std::optional locateProgram(llvm::StringRef envPath, - llvm::StringRef fallbackName) { - if (!envPath.empty() && llvm::sys::fs::exists(envPath)) - return envPath.str(); - if (auto found = llvm::sys::findProgramByName(fallbackName)) - return *found; - return std::nullopt; -} - -class VPTOFatobjToolchain; - -static bool compileDeviceLLVMToObject(llvm::StringRef llPath, - llvm::StringRef outObjPath, - llvm::StringRef targetCPU, - llvm::StringRef bishengPath, - llvm::StringRef stderrPath, - llvm::raw_ostream &diagOS); -static bool compileHostStubToFatobj(llvm::StringRef stubPath, - llvm::StringRef outObjPath, - llvm::StringRef moduleId, - llvm::StringRef targetCPU, - const VPTOFatobjToolchain &toolchain, - llvm::StringRef deviceObjPath, - llvm::StringRef stderrPath, - llvm::raw_ostream &diagOS); -static bool mergeDeviceObjects(llvm::ArrayRef deviceObjPaths, - llvm::StringRef outObjPath, - llvm::StringRef ldLldPath, - llvm::StringRef stderrPath, - llvm::raw_ostream &diagOS); - -class VPTOFatobjToolchain { -public: - static std::optional - create(llvm::raw_ostream &diagOS) { - std::optional ascendHome = getAscendHomePath(); - if (!ascendHome) { - diagOS << "Error: ASCEND_HOME_PATH is required for VPTO fatobj emission.\n"; - return std::nullopt; - } - - VPTOFatobjToolchain toolchain(*ascendHome); - if (!toolchain.validate(diagOS)) - return std::nullopt; - return toolchain; - } - - const std::string &ascendHome() const { return ascendHomePath; } - const std::string &bisheng() const { return bishengPath; } - const std::string &bishengCc1() const { return bishengCc1Path; } - const std::string &cceLd() const { return cceLdPath; } - const std::string &ldLld() const { return ldLldPath; } - const std::string &resourceDir() const { return resourceDirPath; } - const std::string &resourceIncludeDir() const { - return resourceIncludeDirPath; - } - const std::string &cceStubDir() const { return cceStubDirPath; } - const std::string &bishengCompilerBinDir() const { - return bishengCompilerBinDirPath; - } - -private: - explicit VPTOFatobjToolchain(llvm::StringRef ascendHome) - : ascendHomePath(ascendHome.str()), - bishengPath(joinPath(ascendHomePath, "bin/bisheng")), - bishengCc1Path( - joinPath(ascendHomePath, "tools/bisheng_compiler/bin/bisheng")), - cceLdPath(joinPath(ascendHomePath, "bin/cce-ld")), - ldLldPath( - locateProgram(joinPath(ascendHomePath, "bin/ld.lld"), "ld.lld") - .value_or(std::string())), - resourceDirPath(joinPath( - ascendHomePath, "tools/bisheng_compiler/lib/clang/15.0.5")), - resourceIncludeDirPath(joinPath(resourceDirPath, "include")), - cceStubDirPath(joinPath(resourceIncludeDirPath, "cce_stub")), - bishengCompilerBinDirPath( - joinPath(ascendHomePath, "tools/bisheng_compiler/bin")) {} - - bool validate(llvm::raw_ostream &diagOS) const { - if (!llvm::sys::fs::exists(bishengPath)) { - diagOS << "Error: unable to locate bisheng: " << bishengPath << "\n"; - return false; - } - if (!llvm::sys::fs::exists(bishengCc1Path)) { - diagOS << "Error: unable to locate bisheng cc1 frontend: " - << bishengCc1Path << "\n"; - return false; - } - if (!llvm::sys::fs::exists(cceLdPath)) { - diagOS << "Error: unable to locate cce-ld: " << cceLdPath << "\n"; - return false; - } - if (ldLldPath.empty() || !llvm::sys::fs::exists(ldLldPath)) { - diagOS << "Error: unable to locate ld.lld.\n"; - return false; - } - return true; - } - - 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; -}; - -class VPTOFatobjArtifacts { -public: - explicit VPTOFatobjArtifacts(TempFileRegistry &tempFiles) - : tempFiles(tempFiles) {} - - bool emitStubSource(StringRef stubSource, llvm::raw_ostream &diagOS) { - if (!tempFiles.create("ptoas-host-stub", ".cpp", stubPath, diagOS)) - return false; - if (!writeTextFile(stubPath, stubSource, diagOS)) - return false; - return true; - } - - bool initCommandLogs(llvm::raw_ostream &diagOS) { - if (!tempFiles.create("ptoas-stderr", ".log", stderrPath, diagOS)) - return false; - return true; - } - - bool emitCubeObject(llvm::Module *module, - const VPTOFatobjToolchain &toolchain, - llvm::raw_ostream &diagOS) { - if (!module) - return true; - if (!tempFiles.create("ptoas-device", ".ll", cubeLLPath, diagOS)) - return false; - if (!writeLLVMModuleFile(*module, cubeLLPath, diagOS)) - return false; - if (!tempFiles.create("ptoas-device", ".o", cubeObjPath, diagOS)) - return false; - return compileDeviceLLVMToObject(cubeLLPath, cubeObjPath, - "dav-c310-cube", toolchain.bisheng(), - stderrPath, diagOS); - } - - bool emitVectorObject(llvm::Module *module, - const VPTOFatobjToolchain &toolchain, - llvm::raw_ostream &diagOS) { - if (!module) - return true; - if (!tempFiles.create("ptoas-device", ".ll", vectorLLPath, diagOS)) - return false; - if (!writeLLVMModuleFile(*module, vectorLLPath, diagOS)) - return false; - if (!tempFiles.create("ptoas-device", ".o", vectorObjPath, diagOS)) - return false; - return compileDeviceLLVMToObject(vectorLLPath, vectorObjPath, - "dav-c310-vec", toolchain.bisheng(), - stderrPath, diagOS); - } - - bool mergeDeviceObjects(const VPTOFatobjToolchain &toolchain, - llvm::raw_ostream &diagOS) { - llvm::SmallVector deviceObjPaths; - if (!cubeObjPath.empty()) - deviceObjPaths.push_back(cubeObjPath); - if (!vectorObjPath.empty()) - deviceObjPaths.push_back(vectorObjPath); - if (deviceObjPaths.empty()) { - diagOS << "Error: VPTO fatobj emission requires at least one device module.\n"; - return false; - } - if (!tempFiles.create("ptoas-device-merged", ".o", mergedDeviceObjPath, - diagOS)) - return false; - return ::mergeDeviceObjects(deviceObjPaths, mergedDeviceObjPath, - toolchain.ldLld(), stderrPath, diagOS); - } - - bool compileHostStub(const VPTOFatobjToolchain &toolchain, - llvm::StringRef moduleId, - llvm::StringRef targetCPU, - llvm::raw_ostream &diagOS) { - if (!tempFiles.create("ptoas-host-stub", ".o", hostStubObjPath, diagOS)) - return false; - return compileHostStubToFatobj(stubPath, hostStubObjPath, moduleId, - targetCPU, toolchain, mergedDeviceObjPath, - stderrPath, diagOS); - } - - bool repackFatObj(const VPTOFatobjToolchain &toolchain, - llvm::StringRef moduleId, llvm::StringRef targetCPU, - llvm::StringRef outPath, llvm::raw_ostream &diagOS) { - llvm::SmallVector args = { - toolchain.cceLd(), - toolchain.ldLld(), - "-x", - "-cce-lite-bin-module-id", - moduleId.str(), - std::string("-cce-aicore-arch=") + targetCPU.str(), - "-r", - "-o", - outPath.str(), - "-cce-stub-dir", - toolchain.cceStubDir(), - "-cce-install-dir", - toolchain.bishengCompilerBinDir(), - "-cce-inputs-number", - "1", - hostStubObjPath, - }; - return runCommandWithStderr(toolchain.cceLd(), args, stderrPath, diagOS, - "fatobj repack"); - } - -private: - TempFileRegistry &tempFiles; - std::string cubeLLPath; - std::string cubeObjPath; - std::string vectorLLPath; - std::string vectorObjPath; - std::string mergedDeviceObjPath; - std::string stderrPath; - std::string stubPath; - std::string hostStubObjPath; -}; - -static bool runCommandWithStderr(llvm::StringRef program, - llvm::ArrayRef ownedArgs, - llvm::StringRef stderrPath, - llvm::raw_ostream &diagOS, - llvm::StringRef what, - std::optional stdinPath) { - llvm::SmallVector args; - args.reserve(ownedArgs.size()); - for (const std::string &arg : ownedArgs) - args.push_back(arg); - llvm::SmallVector, 3> redirects = { - stdinPath, std::nullopt, stderrPath}; - - std::string execErr; - bool execFailed = false; - int rc = llvm::sys::ExecuteAndWait(program, args, std::nullopt, redirects, 0, - 0, &execErr, &execFailed); - if (!execFailed && rc == 0) - return true; - - diagOS << "Error: " << what << " failed\n"; - diagOS << "Command:"; - for (llvm::StringRef arg : args) - diagOS << " " << arg; - diagOS << "\n"; - if (!execErr.empty()) - diagOS << execErr << "\n"; - if (auto buffer = llvm::MemoryBuffer::getFile(stderrPath)) - diagOS << buffer.get()->getBuffer() << "\n"; - return false; -} - -static bool compileDeviceLLVMToObject(llvm::StringRef llPath, - llvm::StringRef outObjPath, - llvm::StringRef targetCPU, - llvm::StringRef bishengPath, - llvm::StringRef stderrPath, - llvm::raw_ostream &diagOS) { - llvm::SmallVector args = { - bishengPath.str(), - "--target=hiipu64-hisilicon-cce", - std::string("-march=") + targetCPU.str(), - std::string("--cce-aicore-arch=") + targetCPU.str(), - "--cce-aicore-only", - "-O2", - "-mllvm", - "-cce-dyn-kernel-stack-size=true", - "-c", - "-x", - "ir", - "-", - "-o", - outObjPath.str(), - }; - return runCommandWithStderr(bishengPath, args, stderrPath, diagOS, - "device LLVM compilation", llPath); -} - -static bool compileHostStubToFatobj(llvm::StringRef stubPath, - llvm::StringRef outObjPath, - llvm::StringRef moduleId, - llvm::StringRef targetCPU, - const VPTOFatobjToolchain &toolchain, - llvm::StringRef deviceObjPath, - llvm::StringRef stderrPath, - llvm::raw_ostream &diagOS) { - std::string coverageDir = "."; - std::string debugDir = "."; - std::string hostTriple = llvm::sys::getProcessTriple(); - - llvm::SmallVector args = { - toolchain.bishengCc1(), - "-cc1", - "-triple", - hostTriple, - "-target-cpu", - llvm::sys::getHostCPUName().str(), - "-fcce-aicpu-legacy-launch", - "-fcce-is-host", - "-cce-enable-mix", - "-mllvm", - "-enable-mix=true", - "-cce-launch-with-flagv2-impl", - "-fcce-aicore-arch", - targetCPU.str(), - "-fcce-fatobj-compile", - "-emit-obj", - "--mrelax-relocations", - "-disable-free", - "-clear-ast-before-backend", - "-disable-llvm-verifier", - "-discard-value-names", - "-main-file-name", - "stub.cpp", - "-mrelocation-model", - "pic", - "-pic-level", - "2", - "-fhalf-no-semantic-interposition", - "-mframe-pointer=none", - "-fmath-errno", - "-ffp-contract=on", - "-fno-rounding-math", - "-mconstructor-aliases", - "-funwind-tables=2", - "-fallow-half-arguments-and-returns", - "-mllvm", - "-treat-scalable-fixed-error-as-warning", - std::string("-fcoverage-compilation-dir=") + coverageDir, - "-resource-dir", - toolchain.resourceDir(), - "-internal-isystem", - toolchain.resourceIncludeDir(), - "-include", - "__clang_cce_runtime_wrapper.h", - "-D", - "_FORTIFY_SOURCE=2", - "-D", - "REGISTER_BASE", - "-O2", - "-Wno-macro-redefined", - "-Wno-ignored-attributes", - "-std=c++17", - "-fdeprecated-macro", - std::string("-fdebug-compilation-dir=") + debugDir, - "-ferror-limit", - "19", - "-stack-protector", - "2", - "-fno-signed-char", - "-fgnuc-version=4.2.1", - "-fcxx-exceptions", - "-fexceptions", - "-vectorize-loops", - "-vectorize-slp", - "-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", - "-fcce-include-aibinary", - deviceObjPath.str(), - "-fcce-device-module-id", - moduleId.str(), - "-faddrsig", - "-D__GCC_HAVE_DWARF2_CFI_ASM=1", - "-o", - outObjPath.str(), - "-x", - "cce", - stubPath.str(), - }; - return runCommandWithStderr(toolchain.bishengCc1(), args, stderrPath, diagOS, - "host stub compilation"); -} - -static bool mergeDeviceObjects(llvm::ArrayRef deviceObjPaths, - llvm::StringRef outObjPath, - llvm::StringRef ldLldPath, - llvm::StringRef stderrPath, - llvm::raw_ostream &diagOS) { - if (deviceObjPaths.empty()) - return false; - - llvm::SmallVector args = { - ldLldPath.str(), - "-m", - "aicorelinux", - "-Ttext", - "0", - }; - for (const std::string &path : deviceObjPaths) - args.push_back(path); - args.push_back("-o"); - args.push_back(outObjPath.str()); - args.push_back("-r"); - args.push_back("--allow-multiple-definition"); - return runCommandWithStderr(ldLldPath, args, stderrPath, diagOS, - "device object merge"); -} - -} // namespace - -mlir::LogicalResult mlir::pto::emitVPTOFatobj(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(); - } - - std::optional toolchain = - VPTOFatobjToolchain::create(diagOS); - if (!toolchain) - return failure(); - - TempFileRegistry tempFiles; - VPTOFatobjArtifacts artifacts(tempFiles); - if (!artifacts.emitStubSource(stubSource, diagOS)) - return failure(); - if (!artifacts.initCommandLogs(diagOS)) - return failure(); - - if (!artifacts.emitCubeObject(cubeModule, *toolchain, diagOS)) - return failure(); - if (!artifacts.emitVectorObject(vectorModule, *toolchain, diagOS)) - return failure(); - - if (!artifacts.mergeDeviceObjects(*toolchain, diagOS)) - return failure(); - - std::string moduleId = sanitizeModuleId(outputFile.getFilename()); - constexpr llvm::StringLiteral hostTargetCPU = "dav-c310"; - if (!artifacts.compileHostStub(*toolchain, moduleId, hostTargetCPU, diagOS)) - return failure(); - - if (!artifacts.repackFatObj(*toolchain, moduleId, hostTargetCPU, - outputFile.getFilename(), diagOS)) - return failure(); - outputFile.keep(); - return success(); -} diff --git a/tools/ptoas/VPTOHostStubEmission.cpp b/tools/ptoas/VPTOHostStubEmission.cpp index c091d5c1ad..74fde419d2 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 (!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..c954bce626 --- /dev/null +++ b/tools/ptoas/driver.cpp @@ -0,0 +1,723 @@ +// 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 "VPTOHostStubEmission.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/SymbolTable.h" +#include "mlir/Parser/Parser.h" +#include "ptobc/ptobc_decode.h" +#include "llvm/ADT/ScopeExit.h" +#include "llvm/ADT/StringMap.h" +#include "llvm/IR/Function.h" +#include "llvm/IR/GlobalValue.h" +#include "llvm/IR/Module.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/Support/Regex.h" +#include "llvm/Support/SourceMgr.h" +#include "llvm/Support/ToolOutputFile.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 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 std::string normalizePTOASArch(llvm::StringRef archValue) { + std::string normalized = archValue.str(); + for (char &c : normalized) + c = static_cast(std::tolower(static_cast(c))); + return normalized; +} + +static bool isSupportedPTOASArch(llvm::StringRef archValue) { + return archValue == "a3" || archValue == "a5"; +} + +static std::optional +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 isPTOBCBuffer(llvm::StringRef buffer) { + return buffer.size() >= 6 && std::memcmp(buffer.data(), "PTOBC\0", 6) == 0; +} + +static std::unique_ptr readInputBuffer() { + auto fileOrErr = llvm::MemoryBuffer::getFileOrSTDIN(inputFilename); + if (!fileOrErr) { + llvm::errs() << "Error: Could not open input file: " + << fileOrErr.getError().message() << "\n"; + return nullptr; + } + return std::move(*fileOrErr); +} + +static bool resolveTextInputArch(llvm::StringRef buffer, bool cliArchSpecified, + std::string &arch) { + arch = normalizePTOASArch(mlir::pto::ptoTargetArch); + if (cliArchSpecified) { + if (!isSupportedPTOASArch(arch)) { + llvm::errs() << "Error: invalid --pto-arch='" << mlir::pto::ptoTargetArch + << "'. Expected 'a3' or 'a5'.\n"; + return false; + } + return true; + } + + if (auto detectedArch = detectPTOASTextualModuleArch(buffer)) + arch = *detectedArch; + if (!isSupportedPTOASArch(arch)) + arch = "a3"; + return true; +} + +static OwningOpRef decodePTOBCModule(llvm::StringRef buffer, + MLIRContext &context) { + llvm::ArrayRef bytes(reinterpret_cast(buffer.data()), + buffer.size()); +#if defined(__cpp_exceptions) || defined(__EXCEPTIONS) + try { + return ptobc::decodePTOBCToModule(bytes, context); + } catch (...) { + llvm::errs() << "Error: Failed to decode PTOBC.\n"; + return {}; + } +#else + OwningOpRef module = ptobc::decodePTOBCToModule(bytes, context); + if (!module) + llvm::errs() << "Error: Failed to decode PTOBC.\n"; + return module; +#endif +} + +static OwningOpRef +parseTextualModule(std::unique_ptr inputBuffer, + MLIRContext &context, llvm::StringRef arch) { + llvm::SourceMgr sourceMgr; + sourceMgr.AddNewSourceBuffer(std::move(inputBuffer), llvm::SMLoc()); + mlir::pto::ScopedPTOParserTargetArch scopedParserArch( + &context, arch == "a5" ? mlir::pto::PTOParserTargetArch::A5 + : mlir::pto::PTOParserTargetArch::A3); + OwningOpRef module = parseSourceFile(sourceMgr, &context); + if (!module) + llvm::errs() << "Error: Failed to parse MLIR.\n"; + return module; +} + +static OwningOpRef +loadInputModule(std::unique_ptr inputBuffer, + MLIRContext &context, bool cliArchSpecified, + std::string &arch) { + llvm::StringRef buffer = inputBuffer->getBuffer(); + + OwningOpRef module; + if (isPTOBCBuffer(buffer)) { + arch = normalizePTOASArch(mlir::pto::ptoTargetArch); + if (cliArchSpecified && !isSupportedPTOASArch(arch)) { + llvm::errs() << "Error: invalid --pto-arch='" << mlir::pto::ptoTargetArch + << "'. Expected 'a3' or 'a5'.\n"; + return {}; + } + module = decodePTOBCModule(buffer, context); + } else { + if (!resolveTextInputArch(buffer, cliArchSpecified, arch)) + return {}; + module = parseTextualModule(std::move(inputBuffer), context, arch); + } + if (!module) + return {}; + + Operation *moduleOp = module.get().getOperation(); + if (cliArchSpecified) { + moduleOp->setAttr("pto.target_arch", + mlir::StringAttr::get(moduleOp->getContext(), arch)); + return module; + } + + if (auto archAttr = moduleOp->getAttrOfType("pto.target_arch")) { + std::string moduleArch = normalizePTOASArch(archAttr.getValue()); + if (isSupportedPTOASArch(moduleArch)) { + arch = std::move(moduleArch); + return module; + } + } + + if (!isSupportedPTOASArch(arch)) + arch = "a3"; + moduleOp->setAttr("pto.target_arch", + mlir::StringAttr::get(moduleOp->getContext(), arch)); + return module; +} + +static bool parseDriverBackend(llvm::StringRef backendStr, + mlir::pto::PTOBackend &out) { + std::string s = backendStr.str(); + for (char &c : s) + c = static_cast(std::tolower(static_cast(c))); + if (s == "emitc") { + out = mlir::pto::PTOBackend::EmitC; + return true; + } + if (s == "vpto") { + out = mlir::pto::PTOBackend::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::PTOBackend attrBackend = mlir::pto::PTOBackend::EmitC; + if (!parseDriverBackend(backendAttr.getValue(), attrBackend)) { + return op->emitError("invalid pto.backend '") + << backendAttr.getValue() << "'. Expected 'emitc' or 'vpto'."; + } + + backend = attrBackend; + return success(); +} + +static OwningOpRef detachBackendChildModule(ModuleOp outer, + ModuleOp child) { + for (NamedAttribute attr : outer->getAttrs()) { + StringRef attrName = attr.getName().getValue(); + if (attrName == SymbolTable::getSymbolAttrName() || + attrName == "pto.backend") + continue; + if (!child->hasAttr(attr.getName())) + child->setAttr(attr.getName(), attr.getValue()); + } + + child.getOperation()->remove(); + return OwningOpRef(child); +} + +static constexpr llvm::StringLiteral kEmptyHostStubSource = + "#ifndef __global__\n#define __global__\n#endif\n\n" + "#ifndef __gm__\n#define __gm__\n#endif\n\n"; + +class PTOASContext; + +static LogicalResult emitVPTOLLVMFatobj( + const mlir::pto::PTOASCompileResult &jobResult, PTOASContext &context, + llvm::StringRef moduleId, llvm::StringRef outputPath); + +class PTOASContext { +public: + PTOASContext(DialectRegistry ®istry, llvm::StringRef outputPath, int argc, + char **argv) + : mlirContext(registry), outputPath(outputPath.str()), argc(argc), + argv(argv) {} + + ~PTOASContext() = default; + + void initializeMLIRContext() { + // Be tolerant: ptobc decode may materialize ops from dialects that aren't + // explicitly registered/loaded in this tool yet. + mlirContext.allowUnregisteredDialects(true); + mlir::pto::loadPTOASDialects(mlirContext); + } + + LogicalResult initializeObjectEmission() { + return mlir::pto::discoverObjectEmissionToolchain(toolchain, llvm::errs()); + } + + MLIRContext &getMLIRContext() { return mlirContext; } + + void setArch(std::string value) { arch = std::move(value); } + + llvm::StringRef getArch() const { return arch; } + + int getArgc() const { return argc; } + + char **getArgv() const { return argv; } + + llvm::StringRef getOutputPath() const { return outputPath; } + + std::string allocModuleId() { + static size_t nextModuleId = 0; + return "ptoas_module_" + std::to_string(nextModuleId++); + } + + const mlir::pto::ObjectEmissionToolchain &getToolchain() const { + return toolchain; + } + + mlir::pto::TempFileRegistry &getTempFiles() { return tempFiles; } + + LogicalResult createTempPath(llvm::StringRef prefix, llvm::StringRef suffix, + std::string &path) { + return tempFiles.create(prefix, suffix, path, llvm::errs()); + } + +private: + MLIRContext mlirContext; + std::string outputPath; + std::string arch; + int argc = 0; + char **argv = nullptr; + mlir::pto::ObjectEmissionToolchain toolchain; + mlir::pto::TempFileRegistry tempFiles; +}; + +static bool hasPTOKernel(ModuleOp module) { + bool found = false; + module.walk([&](func::FuncOp func) { + if (mlir::pto::isPTOKernelFunction(func)) { + found = true; + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + return found; +} + +class EmitCBackendJob { +public: + EmitCBackendJob(OwningOpRef &module, + mlir::pto::PTOASCompileResult &result) + : module(module), result(result) {} + + LogicalResult run(PTOASContext &context); + +private: + OwningOpRef &module; + mlir::pto::PTOASCompileResult &result; +}; + +class VPTOBackendJob { +public: + VPTOBackendJob(OwningOpRef &module, + mlir::pto::PTOASCompileResult &result) + : module(module), result(result) {} + + LogicalResult run(PTOASContext &context); + +private: + OwningOpRef &module; + mlir::pto::PTOASCompileResult &result; +}; + +class BackendChildJob { +public: + virtual ~BackendChildJob() = default; + virtual LogicalResult run(PTOASContext &context) = 0; +}; + +class EmitCBackendChildJob final : public BackendChildJob { +public: + EmitCBackendChildJob(OwningOpRef &&module, + SmallVectorImpl &fatobjPaths) + : module(std::move(module)), fatobjPaths(fatobjPaths) {} + + LogicalResult run(PTOASContext &context) override { + ModuleOp op = module.get(); + op->setAttr("pto.backend", StringAttr::get(op.getContext(), "emitc")); + + mlir::pto::PTOASCompileResult jobResult; + if (mlir::pto::compilePTOASModule(module, context.getArch(), + mlir::pto::PTOBackend::EmitC, + context.getArgc(), context.getArgv(), + jobResult, + /*emitVPTOHostStub=*/false) != 0) + return failure(); + if (jobResult.kind != mlir::pto::PTOASCompileResultKind::Text) { + llvm::errs() << "Error: EmitC backend child job produced non-text " + "output.\n"; + return failure(); + } + + std::string fatobjPath; + if (failed(context.createTempPath("ptoas-emitc-fatobj", ".o", fatobjPath))) + return failure(); + if (failed(mlir::pto::emitFatobjCCE( + jobResult.textOutput, fatobjPath, context.getToolchain(), + context.getTempFiles(), llvm::errs()))) + return failure(); + + fatobjPaths.push_back(std::move(fatobjPath)); + return success(); + } + +private: + OwningOpRef module; + SmallVectorImpl &fatobjPaths; +}; + +class VPTOBackendChildJob final : public BackendChildJob { +public: + VPTOBackendChildJob(OwningOpRef &&module, std::string moduleId, + SmallVectorImpl &fatobjPaths) + : module(std::move(module)), moduleId(std::move(moduleId)), + fatobjPaths(fatobjPaths) {} + + LogicalResult run(PTOASContext &context) override { + ModuleOp op = module.get(); + op->setAttr("pto.backend", StringAttr::get(op.getContext(), "vpto")); + + bool emitHostStub = hasPTOKernel(op); + mlir::pto::PTOASCompileResult jobResult; + if (mlir::pto::compilePTOASModule( + module, context.getArch(), mlir::pto::PTOBackend::VPTO, + context.getArgc(), context.getArgv(), jobResult, emitHostStub) != 0) + return failure(); + if (jobResult.kind != mlir::pto::PTOASCompileResultKind::VPTOObject) { + llvm::errs() << "Error: VPTO backend child job produced non-object " + "output.\n"; + return failure(); + } + + std::string fatobjPath; + if (failed(context.createTempPath("ptoas-vpto-fatobj", ".o", fatobjPath))) + return failure(); + + if (failed(emitVPTOLLVMFatobj(jobResult, context, moduleId, fatobjPath))) + return failure(); + + fatobjPaths.push_back(std::move(fatobjPath)); + return success(); + } + +private: + OwningOpRef module; + std::string moduleId; + SmallVectorImpl &fatobjPaths; +}; + +class FatobjLinkJob { +public: + explicit FatobjLinkJob(ArrayRef fatobjPaths) + : fatobjPaths(fatobjPaths) {} + + LogicalResult run(PTOASContext &context) { + if (fatobjPaths.size() < 2) { + llvm::errs() + << "Error: mixed backend link requires at least two fatobjs.\n"; + return failure(); + } + + std::string stderrPath; + if (failed(context.createTempPath("ptoas-fatobj", ".log", stderrPath))) + return failure(); + return mlir::pto::linkFatobjs(fatobjPaths, context.getOutputPath(), + context.getToolchain(), stderrPath, + llvm::errs()); + } + +private: + ArrayRef fatobjPaths; +}; + +LogicalResult EmitCBackendJob::run(PTOASContext &context) { + ModuleOp op = module.get(); + op->setAttr("pto.backend", StringAttr::get(op.getContext(), "emitc")); + + if (mlir::pto::compilePTOASModule(module, context.getArch(), + mlir::pto::PTOBackend::EmitC, + context.getArgc(), context.getArgv(), result, + /*emitVPTOHostStub=*/false) != 0) + return failure(); + if (result.kind != mlir::pto::PTOASCompileResultKind::Text) { + llvm::errs() << "Error: EmitC backend job produced non-text output.\n"; + return failure(); + } + return success(); +} + +LogicalResult VPTOBackendJob::run(PTOASContext &context) { + ModuleOp op = module.get(); + op->setAttr("pto.backend", StringAttr::get(op.getContext(), "vpto")); + + bool emitHostStub = hasPTOKernel(op); + if (mlir::pto::compilePTOASModule( + module, context.getArch(), mlir::pto::PTOBackend::VPTO, + context.getArgc(), context.getArgv(), result, emitHostStub) != 0) + return failure(); + if (result.kind == mlir::pto::PTOASCompileResultKind::Text) + return success(); + if (result.kind != mlir::pto::PTOASCompileResultKind::VPTOObject) { + llvm::errs() << "Error: VPTO backend job produced non-VPTO output.\n"; + return failure(); + } + + if (context.getOutputPath().empty() || context.getOutputPath() == "-") { + llvm::errs() << "Error: object output requires an explicit file path " + "passed with -o.\n"; + return failure(); + } + if (failed(context.initializeObjectEmission())) + return failure(); + + std::string moduleId = context.allocModuleId(); + if (failed(emitVPTOLLVMFatobj(result, context, moduleId, + context.getOutputPath()))) + return failure(); + + result.reset(); + result.kind = mlir::pto::PTOASCompileResultKind::MixedObject; + return success(); +} + +static LogicalResult emitVPTOLLVMFatobj( + const mlir::pto::PTOASCompileResult &jobResult, PTOASContext &context, + llvm::StringRef moduleId, llvm::StringRef outputPath) { + llvm::StringRef stubSource = kEmptyHostStubSource; + if (!jobResult.vptoStubSource.empty()) + stubSource = jobResult.vptoStubSource; + + if (failed(mlir::pto::emitFatobjLLVM( + jobResult.vptoCubeModule.module.get(), + jobResult.vptoVectorModule.module.get(), stubSource, + outputPath, moduleId, context.getToolchain(), context.getTempFiles(), + llvm::errs()))) + return failure(); + return success(); +} + +static LogicalResult collectChildJobs( + ModuleOp module, mlir::pto::PTOBackend defaultBackend, + PTOASContext &context, SmallVectorImpl &fatobjPaths, + SmallVectorImpl> &backendJobs) { + SmallVector children(module.getOps()); + for (ModuleOp child : children) { + std::optional childBackend; + if (failed(parseDriverBackendAttr(child.getOperation(), childBackend))) + return failure(); + + OwningOpRef jobModule = detachBackendChildModule(module, child); + if (childBackend.value_or(defaultBackend) == mlir::pto::PTOBackend::VPTO) + backendJobs.push_back(std::make_unique( + std::move(jobModule), context.allocModuleId(), fatobjPaths)); + else + backendJobs.push_back(std::make_unique( + std::move(jobModule), fatobjPaths)); + } + return success(); +} + +static LogicalResult resolveSingleBackend( + bool cliBackendSpecified, + std::optional moduleBackend, + mlir::pto::PTOBackend defaultBackend, ModuleOp module, + std::optional &singleBackend) { + singleBackend = std::nullopt; + if (cliBackendSpecified) + singleBackend = defaultBackend; + if (moduleBackend) + singleBackend = *moduleBackend; + if (singleBackend) + return success(); + + std::optional firstChildBackend; + for (ModuleOp child : module.getOps()) { + std::optional childBackend; + if (failed(parseDriverBackendAttr(child.getOperation(), childBackend))) + return failure(); + + mlir::pto::PTOBackend effectiveChildBackend = + childBackend.value_or(defaultBackend); + if (!firstChildBackend) { + firstChildBackend = effectiveChildBackend; + continue; + } + if (*firstChildBackend != effectiveChildBackend) + return success(); + } + + if (firstChildBackend) + singleBackend = *firstChildBackend; + else + singleBackend = defaultBackend; + return success(); +} + +static LogicalResult runPTOASJobs(OwningOpRef &module, + bool cliBackendSpecified, + PTOASContext &context, + mlir::pto::PTOASCompileResult &result) { + mlir::pto::PTOBackend defaultBackend = mlir::pto::PTOBackend::EmitC; + if (!parseDriverBackend(mlir::pto::ptoBackend, defaultBackend)) { + llvm::errs() << "Error: invalid --pto-backend='" << mlir::pto::ptoBackend + << "'. Expected 'emitc' or 'vpto'.\n"; + return failure(); + } + + std::optional moduleBackend; + if (!cliBackendSpecified) { + if (failed(parseDriverBackendAttr(module->getOperation(), moduleBackend))) + return failure(); + } + + std::optional singleBackend; + if (failed(resolveSingleBackend(cliBackendSpecified, moduleBackend, + defaultBackend, module.get(), singleBackend))) + return failure(); + + if (singleBackend) { + if (*singleBackend == mlir::pto::PTOBackend::EmitC) { + EmitCBackendJob singleJob(module, result); + return singleJob.run(context); + } + VPTOBackendJob singleJob(module, result); + return singleJob.run(context); + } + + SmallVector, 4> backendJobs; + SmallVector fatobjPaths; + if (failed(collectChildJobs(module.get(), defaultBackend, context, fatobjPaths, + backendJobs))) + return failure(); + + if (mlir::pto::emitMlirIR || mlir::pto::emitVPTO || + mlir::pto::ptoPrintSeamIR || !mlir::pto::ptoSeamIRFile.empty()) { + llvm::errs() << "Error: mixed pto.backend fatobj mode does not support " + "debug IR output flags.\n"; + return failure(); + } + if (context.getOutputPath().empty() || context.getOutputPath() == "-") { + llvm::errs() << "Error: mixed pto.backend fatobj mode requires an " + "explicit file path passed with -o.\n"; + return failure(); + } + + result.reset(); + result.kind = mlir::pto::PTOASCompileResultKind::MixedObject; + + if (failed(context.initializeObjectEmission())) + return failure(); + + for (size_t i = 0, e = backendJobs.size(); i < e; ++i) { + if (failed(backendJobs[i]->run(context))) + return failure(); + } + + FatobjLinkJob linkJob(fatobjPaths); + if (failed(linkJob.run(context))) + return failure(); + + return success(); +} + +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(); +} + +// PTOAS driver jobs: +// +----------------------------------------------------------+ +// | .pto | +// +----------------------------------------------------------+ +// +-------------+ +------------+ +------------+ +------------+ +// | EmitC job | | VPTO job | | EmitC | | VPTO | +// | | | | | child job | | child job | +// | | | | +------------+ +------------+ +// | | | | +---------------------------+ +// | | | | | Fatobj link job | +// +-------------+ +------------+ +---------------------------+ +// +-------------+ +------------------------------------------+ +// | C++ source | | fatobj | +// +-------------+ +------------------------------------------+ +int main(int argc, char **argv) { + DialectRegistry registry; + mlir::pto::registerPTOASDialects(registry); + mlir::pto::registerPTOASPassesAndCLOptions(); + 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"); + + std::unique_ptr inputBuffer = readInputBuffer(); + if (!inputBuffer) + return 1; + PTOASContext context(registry, outputFilename, argc, argv); + context.initializeMLIRContext(); + + std::string arch; + OwningOpRef module = loadInputModule( + std::move(inputBuffer), context.getMLIRContext(), cliArchSpecified, arch); + if (!module) + return 1; + context.setArch(std::move(arch)); + + mlir::pto::PTOASCompileResult result; + if (failed(runPTOASJobs(module, cliBackendSpecified, context, result))) + return 1; + + if (result.kind == mlir::pto::PTOASCompileResultKind::Text) + return failed(writeTextOutput(result.textOutput, context.getOutputPath())); + if (result.kind == mlir::pto::PTOASCompileResultKind::MixedObject) + return 0; + + llvm::errs() << "Error: unsupported ptoas compile result.\n"; + return 1; +} diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index caa32e5b36..ffdf965569 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -6,11 +6,11 @@ // 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" @@ -28,7 +28,9 @@ #include #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h" #include "mlir/Target/Cpp/CppEmitter.h" #include "llvm/Support/SourceMgr.h" #include "llvm/Support/ToolOutputFile.h" @@ -72,8 +74,45 @@ using namespace pto; int main(int argc, char **argv); -static void printPTOASVersion(llvm::raw_ostream &os) { - os << "ptoas " << PTOAS_RELEASE_VERSION << "\n"; +void mlir::pto::registerPTOASDialects(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); + pto::registerBufferizableOpInterfaceExternalModels(registry); + + registry.insert(); + registry.insert(); +} + +void mlir::pto::registerPTOASPassesAndCLOptions() { + mlir::registerAllPasses(); + mlir::pto::registerPTOPasses(); + mlir::pto::registerPTOViewToMemrefPass(); + mlir::pto::registerPTOInlineLibCall(); + mlir::pto::registerFoldTileBufIntrinsics(); + mlir::pto::registerExpandTileOp(); + mlir::registerPassManagerCLOptions(); +} + +void mlir::pto::loadPTOASDialects(MLIRContext &context) { + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.getOrLoadDialect(); + context.getOrLoadDialect(); } static std::string getParentDir(llvm::StringRef path) { @@ -256,15 +295,6 @@ static LogicalResult reorderEmitCFunctions(ModuleOp module) { // -------------------------------------------------------------------------- // 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)); @@ -374,12 +404,12 @@ static llvm::cl::opt emitAddPtrTrace( llvm::cl::desc("Emit addptr trace comments in generated C++ output"), llvm::cl::init(false)); -static llvm::cl::opt emitMlirIR( +llvm::cl::opt mlir::pto::emitMlirIR( "emit-pto-ir", llvm::cl::desc("Emit PTO IR after lowering instead of C++"), llvm::cl::init(false)); -static llvm::cl::opt ptoTargetArch( +llvm::cl::opt mlir::pto::ptoTargetArch( "pto-arch", llvm::cl::desc("Target Ascend architecture for codegen: a3 or a5 (default: a3)"), llvm::cl::value_desc("a3|a5"), @@ -391,12 +421,12 @@ static llvm::cl::opt ptoBuildLevel( llvm::cl::value_desc("level1|level2|level3"), llvm::cl::init("level2")); -static llvm::cl::opt ptoBackend( +llvm::cl::opt mlir::pto::ptoBackend( "pto-backend", llvm::cl::desc("Final PTOAS backend: emitc or vpto (default: emitc)"), llvm::cl::value_desc("emitc|vpto"), llvm::cl::init("emitc")); -static llvm::cl::opt emitVPTO( +llvm::cl::opt mlir::pto::emitVPTO( "emit-vpto", llvm::cl::desc("Write final post-pass VPTO IR to -o"), llvm::cl::init(false)); @@ -417,12 +447,12 @@ static llvm::cl::opt dumpVPTOIR( llvm::cl::desc("Print post-pass VPTO backend IR to stderr"), llvm::cl::init(false)); -static llvm::cl::opt ptoPrintSeamIR( +llvm::cl::opt mlir::pto::ptoPrintSeamIR( "pto-print-seam-ir", llvm::cl::desc("Print shared pre-backend seam IR to stderr"), llvm::cl::init(false)); -static llvm::cl::opt ptoSeamIRFile( +llvm::cl::opt mlir::pto::ptoSeamIRFile( "pto-seam-ir-file", llvm::cl::desc("Write shared pre-backend seam IR to a file"), llvm::cl::value_desc("path"), @@ -434,11 +464,6 @@ enum class PTOBuildLevel { Level3, }; -enum class PTOBackend { - EmitC, - VPTO, -}; - static PTOBuildLevel defaultBuildLevel() { return PTOBuildLevel::Level2; } @@ -484,21 +509,6 @@ static bool parseAutoSyncTailHint(llvm::StringRef hintStr, std::string &normaliz return false; } -static bool parseBackend(llvm::StringRef backendStr, PTOBackend &out) { - std::string s = backendStr.str(); - for (char &c : s) - c = static_cast(std::tolower(static_cast(c))); - if (s == "emitc") { - out = PTOBackend::EmitC; - return true; - } - if (s == "vpto") { - out = PTOBackend::VPTO; - return true; - } - return false; -} - static LogicalResult emitSharedPreBackendSeamIR(ModuleOp module, llvm::StringRef outputPath) { if (outputPath.empty()) @@ -1233,38 +1243,37 @@ static pto::VPTOEmissionOptions buildVPTOEmissionOptions() { return options; } -static int emitVPTOBackendResult(ModuleOp module, - llvm::ToolOutputFile &outputFile) { +static int emitVPTOBackendResult(ModuleOp module, PTOASCompileResult &result, + bool emitHostStub) { if (emitVPTO) { - module.print(outputFile.os()); - outputFile.os() << "\n"; - outputFile.keep(); + result.kind = PTOASCompileResultKind::Text; + llvm::raw_string_ostream os(result.textOutput); + module.print(os); + os << "\n"; + os.flush(); return 0; } pto::VPTOEmissionOptions options = buildVPTOEmissionOptions(); std::string stubSource; - if (failed(pto::emitVPTOHostStubSource(module, stubSource, llvm::errs()))) { - llvm::errs() << "Error: Failed to emit VPTO host stub source.\n"; - return 1; + if (emitHostStub) { + if (failed(pto::emitVPTOHostStubSource(module, stubSource, llvm::errs()))) { + llvm::errs() << "Error: Failed to emit VPTO host stub source.\n"; + return 1; + } } - pto::EmittedLLVMModule cubeModule; - pto::EmittedLLVMModule vectorModule; if (failed( - pto::lowerVPTOModuleToLLVMModules(module, options, cubeModule, - vectorModule, llvm::errs()))) { + pto::lowerVPTOModuleToLLVMModules(module, options, + result.vptoCubeModule, + result.vptoVectorModule, + llvm::errs()))) { llvm::errs() << "Error: Failed to lower VPTO to LLVM modules.\n"; return 1; } - if (failed(pto::emitVPTOFatobj(cubeModule.module.get(), - vectorModule.module.get(), stubSource, - outputFile, llvm::errs()))) { - llvm::errs() << "Error: Failed to emit VPTO fatobj.\n"; - return 1; - } - outputFile.keep(); + result.vptoStubSource = std::move(stubSource); + result.kind = PTOASCompileResultKind::VPTOObject; return 0; } @@ -1291,153 +1300,19 @@ static LogicalResult runVPTOBackendPipeline(OwningOpRef &module, 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; - if (!parseBackend(ptoBackend, effectiveBackend)) { - llvm::errs() << "Error: invalid --pto-backend='" << ptoBackend - << "'. Expected 'emitc' or 'vpto'.\n"; - return 1; - } +int mlir::pto::compilePTOASModule( + OwningOpRef &module, llvm::StringRef arch, + PTOBackend effectiveBackend, int argc, char **argv, + PTOASCompileResult &result, bool emitVPTOHostStub) { + result.reset(); if (effectiveBackend != PTOBackend::VPTO && (emitVPTO || ptoPrintSeamIR || !ptoSeamIRFile.empty())) { llvm::errs() << "Error: VPTO-specific flags require " - "--pto-backend=vpto.\n"; + "--pto-backend=vpto or pto.backend = \"vpto\".\n"; return 1; } - // Read whole input first (so we can auto-detect .ptobc by magic). - auto fileOrErr = llvm::MemoryBuffer::getFileOrSTDIN(inputFilename); - if (!fileOrErr) { - llvm::errs() << "Error: Could not open input file: " - << fileOrErr.getError().message() << "\n"; - return 1; - } - - MLIRContext context(registry); - // Be tolerant: ptobc decode may materialize ops from dialects that aren't - // explicitly registered/loaded in this tool yet. - context.allowUnregisteredDialects(true); - - context.getOrLoadDialect(); - context.getOrLoadDialect(); - context.getOrLoadDialect(); - context.getOrLoadDialect(); - context.getOrLoadDialect(); - context.getOrLoadDialect(); - context.getOrLoadDialect(); - context.getOrLoadDialect(); - - OwningOpRef module; - llvm::StringRef buf = (*fileOrErr)->getBuffer(); - const bool isPTOBC = (buf.size() >= 6 && std::memcmp(buf.data(), "PTOBC\0", 6) == 0); - - auto normalizeArch = [](llvm::StringRef archValue) { - std::string normalized = archValue.str(); - for (char &c : normalized) - c = static_cast(std::tolower(static_cast(c))); - return normalized; - }; - auto detectTextualModuleArch = [&](llvm::StringRef text) -> std::optional { - llvm::SmallVector matches; - llvm::Regex archRegex( - R"ptoarch("?(pto\.target_arch)"?[[:space:]]*=[[:space:]]*"([[:alpha:][:digit:]_]+)")ptoarch"); - if (!archRegex.match(text, &matches) || matches.size() < 3) - return std::nullopt; - return normalizeArch(matches[2]); - }; - - std::string arch = normalizeArch(ptoTargetArch); - if (cliArchSpecified) { - if (arch != "a3" && arch != "a5") { - llvm::errs() << "Error: invalid --pto-arch='" << ptoTargetArch - << "'. Expected 'a3' or 'a5'.\n"; - return 1; - } - } else if (!isPTOBC) { - if (auto detectedArch = detectTextualModuleArch(buf)) - arch = *detectedArch; - } - if (arch != "a3" && arch != "a5") - arch = "a3"; - - if (isPTOBC) { - // Decode PTO bytecode directly into an MLIR module. - llvm::ArrayRef bytes(reinterpret_cast(buf.data()), buf.size()); -#if defined(__cpp_exceptions) || defined(__EXCEPTIONS) - try { - module = ptobc::decodePTOBCToModule(bytes, context); - } catch (...) { - llvm::errs() << "Error: Failed to decode PTOBC.\n"; - return 1; - } -#else - module = ptobc::decodePTOBCToModule(bytes, context); -#endif - if (!module) { - llvm::errs() << "Error: Failed to decode PTOBC.\n"; - return 1; - } - } else { - // Parse textual MLIR (.pto). - llvm::SourceMgr sourceMgr; - sourceMgr.AddNewSourceBuffer(std::move(*fileOrErr), llvm::SMLoc()); - pto::ScopedPTOParserTargetArch scopedParserArch( - &context, arch == "a5" ? pto::PTOParserTargetArch::A5 - : pto::PTOParserTargetArch::A3); - module = parseSourceFile(sourceMgr, &context); - if (!module) { - llvm::errs() << "Error: Failed to parse MLIR.\n"; - return 1; - } - } - - // If the CLI explicitly requested an arch, it overrides the input module. - // Otherwise, preserve the textual module's arch when present and only fall - // back to the effective default. - if (cliArchSpecified || !module->getOperation()->hasAttr("pto.target_arch")) { - module->getOperation()->setAttr("pto.target_arch", - mlir::StringAttr::get(&context, arch)); - } - PTOBuildLevel effectiveLevel = defaultBuildLevel(); if (!parseBuildLevel(ptoBuildLevel, effectiveLevel)) { llvm::errs() << "Error: invalid --pto-level='" << ptoBuildLevel @@ -1462,7 +1337,7 @@ 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; @@ -1525,14 +1400,6 @@ int main(int argc, char **argv) { return 1; } - // [Fix] ToolOutputFile Usage - std::error_code ec; - llvm::ToolOutputFile outputFile(outputFilename, ec, llvm::sys::fs::OF_None); - if (ec) { - llvm::errs() << ec.message() << "\n"; - return 1; - } - const bool hasTileOpsToExpand = hasUnexpandedTileOps(*module); const bool hasTilelangHelpers = hasTilelangInlineHelpers(*module); @@ -1545,11 +1412,11 @@ int main(int argc, char **argv) { if (failed(runVPTOBackendPipeline(module, argc, argv, hasTileOpsToExpand, hasTilelangHelpers))) return 1; - return emitVPTOBackendResult(module.get(), outputFile); + return emitVPTOBackendResult(*module, result, emitVPTOHostStub); } // Main PassManager - PassManager pm(&context); + PassManager pm(module->getContext()); if (failed(applyPassManagerCLOptions(pm))) return 1; @@ -1591,15 +1458,15 @@ int main(int argc, char **argv) { pto::createPTOGraphSyncSolverPass(graphSyncOpts)); } - llvm::raw_ostream *outputOS = &outputFile.os(); - if (emitMlirIR) { if (failed(pm.run(*module))) { llvm::errs() << "Error: Pass execution failed.\n"; return 1; } - module->print(*outputOS); - outputFile.keep(); + result.kind = PTOASCompileResultKind::Text; + llvm::raw_string_ostream os(result.textOutput); + module->print(os); + os.flush(); return 0; } @@ -1611,7 +1478,7 @@ int main(int argc, char **argv) { return 1; module->getOperation()->setAttr("pto.target_arch", - mlir::StringAttr::get(&context, arch)); + mlir::StringAttr::get(module->getContext(), arch)); if (effectiveBackend == PTOBackend::VPTO) { if (failed(pm.run(*module))) { @@ -1629,7 +1496,7 @@ int main(int argc, char **argv) { if (failed(runVPTOBackendPipeline(module, argc, argv, hasTileOpsToExpand, hasTilelangHelpers))) return 1; - return emitVPTOBackendResult(module.get(), outputFile); + return emitVPTOBackendResult(*module, result, emitVPTOHostStub); } if (arch == "a3") { @@ -1673,10 +1540,7 @@ int main(int argc, char **argv) { rewriteScalarConstantDecls(cppOutput); rewriteHoistedGlobalTensorDecls(cppOutput); - *outputOS << cppOutput; - outputOS->flush(); - - outputFile.keep(); // Success, keep the file - + result.kind = PTOASCompileResultKind::Text; + result.textOutput = std::move(cppOutput); return 0; } diff --git a/tools/ptoas/ptoas.h b/tools/ptoas/ptoas.h new file mode 100644 index 0000000000..f21bcc93ea --- /dev/null +++ b/tools/ptoas/ptoas.h @@ -0,0 +1,70 @@ +// 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 "PTO/Transforms/VPTOLLVMEmitter.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Support/LogicalResult.h" +#include "llvm/Support/CommandLine.h" +#include +#include + +namespace mlir { +class DialectRegistry; +class MLIRContext; +} // namespace mlir + +namespace mlir::pto { + +extern llvm::cl::opt emitMlirIR; +extern llvm::cl::opt ptoTargetArch; +extern llvm::cl::opt ptoBackend; +extern llvm::cl::opt emitVPTO; +extern llvm::cl::opt ptoPrintSeamIR; +extern llvm::cl::opt ptoSeamIRFile; + +enum class PTOBackend { + EmitC, + VPTO, +}; + +enum class PTOASCompileResultKind { + Text, + VPTOObject, + MixedObject, +}; + +struct PTOASCompileResult { + void reset() { + textOutput.clear(); + vptoStubSource.clear(); + vptoCubeModule.reset(); + vptoVectorModule.reset(); + kind = PTOASCompileResultKind::Text; + } + + PTOASCompileResultKind kind = PTOASCompileResultKind::Text; + std::string textOutput; + std::string vptoStubSource; + EmittedLLVMModule vptoCubeModule; + EmittedLLVMModule vptoVectorModule; +}; + +int compilePTOASModule(OwningOpRef &module, llvm::StringRef arch, + PTOBackend backend, int argc, char **argv, + PTOASCompileResult &result, + bool emitVPTOHostStub = true); +void registerPTOASDialects(DialectRegistry ®istry); +void registerPTOASPassesAndCLOptions(); +void loadPTOASDialects(MLIRContext &context); + +} // namespace mlir::pto + +#endif