diff --git a/docker/test_wheel_imports.sh b/docker/test_wheel_imports.sh index 85a85aa143..4c1a0a318f 100755 --- a/docker/test_wheel_imports.sh +++ b/docker/test_wheel_imports.sh @@ -222,13 +222,4 @@ grep -q "candidates = " "${CLEAN_ENV_PTO_IR}" || { echo "Error: clean-environment ptoas smoke output is missing TileLib candidate metadata" >&2 exit 1 } -if ! grep -q "TileLib daemon started successfully" "${CLEAN_ENV_LOG}"; then - echo "Error: TileLib daemon did not report a successful start" >&2 - exit 1 -fi -if ! grep -q "TileLib daemon stopped" "${CLEAN_ENV_LOG}"; then - echo "Error: TileLib daemon did not report a clean stop" >&2 - exit 1 -fi - echo "All wheel import tests passed!" diff --git a/docs/designs/ptoas-compiler-dso-wheel-linking.md b/docs/designs/ptoas-compiler-dso-wheel-linking.md index bdd3bed827..79502a0c35 100644 --- a/docs/designs/ptoas-compiler-dso-wheel-linking.md +++ b/docs/designs/ptoas-compiler-dso-wheel-linking.md @@ -109,8 +109,7 @@ _core │ ├── ptoas.cpp │ ├── driver.cpp │ ├── VPTOHostStubEmission.cpp -│ ├── ObjectEmission.cpp -│ └── TilelangDaemon.cpp +│ └── ObjectEmission.cpp ├── PTOCAPI └── PTOASPythonCAPI ``` diff --git a/docs/designs/ptoas-python-launcher-layout.md b/docs/designs/ptoas-python-launcher-layout.md index 1a9d3d2d29..950ed1b070 100644 --- a/docs/designs/ptoas-python-launcher-layout.md +++ b/docs/designs/ptoas-python-launcher-layout.md @@ -139,7 +139,7 @@ assemble Python packages from unrelated build directories. The current archive is built against CPython 3.11 and requires a CPython 3.11 interpreter. `bin/ptoas` adds the archive root to `sys.path`, then uses the same `ptoas._cli -> ptoas._core` path as the install tree. The packaged `ptodsl/` -tree supports the compiler's default PTODSL TileLib backend; it does not turn +tree supports the compiler's PTODSL TileLib implementation; it does not turn the archive into a normal pip-installable PTODSL distribution. Linux archives use package-relative and archive-relative `$ORIGIN` RPATHs; @@ -155,8 +155,10 @@ declared installation layout. CTest and direct developer-tree runs must set an explicit matching `PYTHONPATH`; PTODSL does not guess repository, LLVM build, or PTOAS install paths at import time. -TileOp expansion remains a lazy, separate daemon process. The PTOAS CLI passes -the packaged PTODSL root and the active Python executable to the native driver, -which starts the daemon only when expansion is required. Keeping the daemon out -of the compiler process also prevents independently packaged MLIR/LLVM Python -bindings from registering runtime state in the native PTOAS process. +TileOp expansion runs in the CLI's existing Python process. `_core.main` creates +one Python-owned MLIR context for the compilation session, and the native driver +borrows that exact context. The in-process TileLib service materializes a source +module in the shared context and clones it into native ownership before the +Python module owner can be released. The packaged MLIR bindings and PTOAS must +therefore remain one ABI-matched runtime rather than independently replaceable +components. diff --git a/docs/designs/ptodsl-tilelib-template-selection-design.md b/docs/designs/ptodsl-tilelib-template-selection-design.md index 16e8b7f503..35d795d651 100644 --- a/docs/designs/ptodsl-tilelib-template-selection-design.md +++ b/docs/designs/ptodsl-tilelib-template-selection-design.md @@ -2,10 +2,7 @@ ## Background -PTOAS currently supports two TileLib backends for VPTO tile-op expansion: - -- `tilelang`, the legacy TileLangDSL template implementation. -- `ptodsl`, the PTODSL-native template implementation. +PTOAS uses the PTODSL-native TileLib implementation for VPTO tile-op expansion. A tile op may have several legal implementations for the same op name. Those implementations can differ by dtype, layout, memory space, @@ -42,7 +39,7 @@ in the ISA and user guide documents, not here. ## Pipeline -The PTODSL TileLib path has two compiler interactions with the Python daemon. +The PTODSL TileLib path has two interactions with the in-process Python service. ```text TileOp in MLIR @@ -50,7 +47,7 @@ TileOp in MLIR | InsertTemplateAttributes | - reconstruct operand specs from MLIR | - collect context attributes - | - ask the PTODSL daemon for legal candidates + | - ask the PTODSL service for legal candidates | - store compact candidate metadata on the TileOp v TileOp with candidates attr @@ -58,8 +55,8 @@ TileOp with candidates attr | ExpandTileOp | - build a specialization key from current MLIR operands and attrs | - choose candidate 0 from the compact candidates attr - | - ask the daemon to render that candidate - | - clone the generated helper and replace the TileOp with func.call + | - ask the service to materialize that candidate in the shared context + | - import the generated entry/helpers and replace the TileOp with func.call v VPTO-facing IR ``` @@ -69,6 +66,20 @@ before later passes can make candidate information harder to reconstruct. `ExpandTileOp` still renders from the current MLIR operands so the helper body matches the actual operand types and view metadata that survived to expansion. +Both stages are ordinary registered MLIR passes. They are default-constructible +and obtain the current `MLIRContext` from the operation being transformed. A +process-wide `TileLibRuntime` provides the host `TileLibService`; it owns no +compilation context and receives the current context explicitly for every +materialization. `PTOASContext` continues to own or borrow the context for one +compilation session, so different invocations may use different contexts while +sharing one Python runtime. + +The Python entry keeps the corresponding Python `Context` owner alive for the +complete native compilation call. Compiler materialization requires that +explicit context and never falls back to creating another one. This preserves +normal pass registration, textual pipelines, targeted IR printing, cloning, +and reproducer behavior without storing Python objects in pass instances. + ## Template Metadata PTODSL template authors register versions through `tilelib.tile_template`. @@ -102,7 +113,7 @@ remain in Python metadata for selection, diagnostics, and future tooling. ## Operand Specs Both `InsertTemplateAttributes` and `ExpandTileOp` reconstruct operand specs -from MLIR. The JSON shape sent to the daemon is deliberately close to +from MLIR. The JSON shape sent to the Python service is deliberately close to `TileSpec`, `ViewSpec`, `ScalarSpec`, and `VectorSpec`. | Operand kind | Required metadata | @@ -137,7 +148,7 @@ template is considered ported. ## Candidate Legality And Ranking -The daemon loads only the template module for the requested op and target. It +The service loads only the template module for the requested op and target. It then evaluates each registered candidate: 1. Bind positional MLIR operands to the template parameter names. @@ -149,7 +160,7 @@ then evaluates each registered candidate: 7. Run custom constraint predicates. 8. Sort legal candidates by descending priority. -If no candidate is legal, the daemon reports a `NoMatchingTemplate` error with +If no candidate is legal, the service reports a `NoMatchingTemplate` error with per-candidate reasons. If multiple candidates tie for the highest priority and no explicit candidate is requested, the registry reports ambiguity rather than silently picking one. @@ -170,7 +181,7 @@ TileOp. Each entry contains: - `tail` This attribute is intentionally not a copy of the full Python metadata object. -Legality has already happened in the daemon. The IR only needs a stable list of +Legality has already happened in the service. The IR only needs a stable list of legal render targets and the small amount of metadata consumed by downstream passes. @@ -180,8 +191,9 @@ Python metadata. Add a field only when a C++ pass or IR-level test consumes it. ## Expansion And Specialization `ExpandTileOp` uses the first candidate in the compact candidate list. For -PTODSL, it passes the selected candidate name back to the daemon so rendering -cannot accidentally choose a different legal template after the metadata pass. +PTODSL, it passes the selected candidate name back to the service so +materialization cannot accidentally choose a different legal template after the +metadata pass. The specialization key deduplicates generated helpers inside one module. It must include every input that can change the rendered helper body: diff --git a/include/PTO/Transforms/Passes.h b/include/PTO/Transforms/Passes.h index d3c02994f9..dca36c2904 100644 --- a/include/PTO/Transforms/Passes.h +++ b/include/PTO/Transforms/Passes.h @@ -23,6 +23,7 @@ #include "llvm/ADT/StringRef.h" #include "mlir/Pass/Pass.h" #include "PTO/IR/PTODialect.h" +#include "PTO/Transforms/TileLibService.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/Arith/IR/Arith.h" @@ -131,10 +132,7 @@ std::unique_ptr createVMILowerUnifiedToLegacyPass(); std::unique_ptr createVMINormalizeSignlessIntToUnsignedPass(); std::unique_ptr createVMIToVPTOPass(); std::unique_ptr createInsertTemplateAttributesPass(); -std::unique_ptr createInsertTemplateAttributesPass( - const InsertTemplateAttributesOptions &options); std::unique_ptr createExpandTileOpPass(); -std::unique_ptr createExpandTileOpPass(const ExpandTileOpOptions &options); std::unique_ptr createFoldTileBufIntrinsicsPass(); std::unique_ptr createFoldTileBufIntrinsicsPass(llvm::StringRef foldMode); std::unique_ptr createPTOCanonicalizeIRPass(); diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index c2f03ff953..251ffbf9b1 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -511,40 +511,26 @@ def InsertTemplateAttributes : Pass<"pto-insert-template-attributes", "ModuleOp"> { let summary = "Attach legal PTODSL template candidates to tile operations"; let description = [{ - Queries the PTODSL TileLib daemon for legal template candidates and stores - the compact candidate list on each tile operation as the `candidates` - attribute. Each candidate contains only id, name, loop_depth, postupdate, - and tail metadata. + Queries the process-wide PTODSL TileLib runtime for legal template + candidates and stores the compact candidate list on each tile operation as + the `candidates` attribute. Each candidate contains only id, name, + loop_depth, postupdate, and tail metadata. }]; let constructor = "mlir::pto::createInsertTemplateAttributesPass()"; let dependentDialects = [ "mlir::pto::PTODialect", "mlir::func::FuncDialect" ]; - let options = [ - Option<"pythonExe", "python-exe", "std::string", - /*default=*/"\"python3\"", - "Python executable for TileLib metadata invocation">, - Option<"daemonSocketPath", "daemon-socket-path", "std::string", - /*default=*/"\"\"", - "Path to the PTODSL TileLib daemon Unix socket">, - Option<"tileLibPkgPath", "tile-lib-pkg-path", "std::string", - /*default=*/"\"\"", - "PYTHONPATH root for PTODSL">, - Option<"daemonHelperModule", "daemon-helper-module", "std::string", - /*default=*/"\"ptodsl.tilelib.serving.helper\"", - "Python module used for daemon metadata RPC calls"> - ]; } def ExpandTileOp : Pass<"pto-expand-tile-op", "ModuleOp"> { let summary = "Expand tile ops into calls to TileLib template functions"; let description = [{ - Expands tile-level operations (pto.tadd, pto.tsub, etc.) by invoking the - selected Python TileLib backend to instantiate template libraries. The - generated template functions use tile_buf parameters and contain - vector-level implementations (pto.vecscope, pto.vlds, pto.vadd, - pto.vsts, etc.). + Expands tile-level operations (pto.tadd, pto.tsub, etc.) by asking the + process-wide PTODSL TileLib runtime to instantiate template libraries in + the current operation's MLIRContext. The generated template functions use + tile_buf parameters and contain vector-level implementations (pto.vecscope, + pto.vlds, pto.vadd, pto.vsts, etc.). Each tile op is replaced by a func.call to the generated template function, with tile_buf operands passed directly (no type bridging). @@ -562,29 +548,12 @@ def ExpandTileOp : Pass<"pto-expand-tile-op", "ModuleOp"> { "mlir::scf::SCFDialect", "mlir::vector::VectorDialect" ]; - let options = [ - Option<"pythonExe", "python-exe", "std::string", - /*default=*/"\"python3\"", - "Python executable for TileLib invocation">, - Option<"daemonSocketPath", "daemon-socket-path", "std::string", - /*default=*/"\"\"", - "Path to Unix domain socket for daemon RPC">, - Option<"tileLibBackend", "tile-lib-backend", "std::string", - /*default=*/"\"ptodsl\"", - "TileLib backend: ptodsl">, - Option<"tileLibPkgPath", "tile-lib-pkg-path", "std::string", - /*default=*/"\"\"", - "PYTHONPATH root for the selected TileLib backend">, - Option<"daemonHelperModule", "daemon-helper-module", "std::string", - /*default=*/"\"ptodsl.tilelib.serving.helper\"", - "Python module used for daemon helper RPC calls"> - ]; } def FoldTileBufIntrinsics : Pass<"pto-fold-tile-buf-intrinsics", "mlir::func::FuncOp"> { let summary = "Fold structured-view intrinsics after template inlining"; let description = [{ - After TileLang DSL template functions are inlined, the IR contains + After PTODSL template functions are inlined, the IR contains structured-view intrinsics whose operands are now bound to concrete values. This pass resolves them: diff --git a/include/PTO/Transforms/TileLibService.h b/include/PTO/Transforms/TileLibService.h new file mode 100644 index 0000000000..a299cd7cd6 --- /dev/null +++ b/include/PTO/Transforms/TileLibService.h @@ -0,0 +1,68 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#ifndef MLIR_DIALECT_PTO_TRANSFORMS_TILELIBSERVICE_H +#define MLIR_DIALECT_PTO_TRANSFORMS_TILELIBSERVICE_H + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Support/LogicalResult.h" + +#include "llvm/ADT/STLFunctionalExtras.h" +#include "llvm/ADT/StringRef.h" + +#include +#include + +namespace mlir::pto { + +/// Pure-data request used by the in-process TileLib materializer. The JSON +/// fields are request data only; generated MLIR never crosses this interface as +/// text. Keeping this interface independent of pybind11 allows transform passes +/// to remain usable from native tests and non-Python hosts. +struct TileLibMaterializationRequest { + std::string target; + std::string op; + std::string operandSpecsJson; + std::string contextAttrsJson; + std::string candidateId; +}; + +using TileLibMaterializationCallback = + llvm::function_ref; + +/// Synchronous handoff for a materialized TileLib implementation. The source +/// module is borrowed and remains owned by the service for the duration of the +/// callback. Consumers must clone/import any operations they need before the +/// callback returns. +class TileLibService { +public: + virtual ~TileLibService() = default; + + virtual FailureOr + getMetadata(const TileLibMaterializationRequest &request) = 0; + + virtual LogicalResult + materialize(const TileLibMaterializationRequest &request, + MLIRContext &context, + TileLibMaterializationCallback callback) = 0; +}; + +/// Process-wide access to the host TileLib implementation. The runtime owns no +/// compilation context: passes obtain the current MLIRContext from their +/// operation and pass it to TileLibService::materialize. A host binding installs +/// one service implementation for the lifetime of that runtime. +class TileLibRuntime { +public: + static void install(std::shared_ptr service); + static void uninstall(TileLibService *service); + static std::shared_ptr getService(); +}; + +} // namespace mlir::pto + +#endif // MLIR_DIALECT_PTO_TRANSFORMS_TILELIBSERVICE_H diff --git a/lib/PTO/Transforms/CMakeLists.txt b/lib/PTO/Transforms/CMakeLists.txt index 61ea155d31..cddb500464 100644 --- a/lib/PTO/Transforms/CMakeLists.txt +++ b/lib/PTO/Transforms/CMakeLists.txt @@ -76,6 +76,7 @@ add_mlir_dialect_library(PTOTransforms InsertSync/InsertSyncDebug.cpp PTORematerializeFixpipeVectorQuant.cpp PTOValidateIntToPtrUses.cpp + TileLibService.cpp InsertTemplateAttributes.cpp ExpandTileOp.cpp FoldTileBufIntrinsics.cpp diff --git a/lib/PTO/Transforms/ExpandTileOp.cpp b/lib/PTO/Transforms/ExpandTileOp.cpp index 95c2a8d511..abc5effb3f 100644 --- a/lib/PTO/Transforms/ExpandTileOp.cpp +++ b/lib/PTO/Transforms/ExpandTileOp.cpp @@ -9,8 +9,8 @@ //===- ExpandTileOp.cpp ---------------------------------------------------===// //===----------------------------------------------------------------------===// // -// Expand tile-level ops (pto.tadd, pto.tsub, ...) by invoking the selected -// Python TileLib backend to instantiate template libraries. +// Expand tile-level ops (pto.tadd, pto.tsub, ...) by materializing PTODSL +// template libraries in the compiler's host Python interpreter. // // The generated template functions use tile_buf parameters. After this pass, // the Inline pass inlines the template body, and FoldTileBufIntrinsics @@ -20,18 +20,18 @@ // 1. Extract SpecKey from ALL operands' tile_buf types. // 2. For PTODSL, read candidates attached by InsertTemplateAttributes and // select the first candidate still present. -// 3. Invoke the selected TileLib helper to generate a specialized MLIR -// function (with tile_buf parameters). -// 4. Parse the generated MLIR and clone the function into the module. +// 3. Ask the in-process TileLib service to build a source module in the same +// MLIRContext. +// 4. Clone its entry/helper functions into the caller module. // 5. Replace the original tile op with func.call, passing tile_buf // operands directly (no type bridging needed). // #include "PTO/IR/PTO.h" #include "PTO/IR/PTOTypeUtils.h" -#include "PTO/Support/PythonExecutable.h" #include "PTO/Transforms/Passes.h" #include "PTO/Transforms/TileOpExpansionUtils.h" +#include "PTO/Transforms/TileLibService.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Func/IR/FuncOps.h" @@ -44,7 +44,6 @@ #include "mlir/IR/IRMapping.h" #include "mlir/IR/SymbolTable.h" #include "mlir/Pass/Pass.h" -#include "mlir/Parser/Parser.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/SmallVector.h" @@ -53,20 +52,10 @@ #include "llvm/ADT/StringSet.h" #include "llvm/ADT/StringExtras.h" #include "llvm/ADT/StringRef.h" -#include "llvm/Support/FileSystem.h" -#include "llvm/Support/MemoryBuffer.h" -#include "llvm/Support/Path.h" -#include "llvm/Support/Program.h" #include "llvm/Support/raw_ostream.h" -#include #include #include -#include - -extern "C" { -extern char **environ; -} using namespace mlir; @@ -74,8 +63,8 @@ namespace mlir { namespace pto { namespace func = ::mlir::func; - #define GEN_PASS_DEF_EXPANDTILEOP - #include "PTO/Transforms/Passes.h.inc" +#define GEN_PASS_DEF_EXPANDTILEOP +#include "PTO/Transforms/Passes.h.inc" } // namespace pto } // namespace mlir @@ -789,19 +778,13 @@ static std::optional buildSpecKey(Operation *op) { // ExpandState: runtime state for a single pass invocation. // ============================================================================ struct ExpandState { - std::vector> parsedModules; // Keep parsed modules alive + std::shared_ptr tileLibService; - std::string tileLibPkgPath; - std::string daemonHelperModule; - std::string pythonExe; - std::string daemonSocketPath; - - std::optional - invokeTileLibHelper(const SpecKey &key, StringRef candidateId = {}); func::FuncOp invokeTileLib(const SpecKey &key, Operation *tileOp, ModuleOp mod, MLIRContext *ctx); - func::FuncOp invokeTileLibDaemon(const SpecKey &key, StringRef candidateId, - ModuleOp mod, MLIRContext *ctx); + func::FuncOp invokeInProcessTileLib(const SpecKey &key, + StringRef candidateId, ModuleOp mod, + MLIRContext *ctx); LogicalResult expandTileOpsInFunction(func::FuncOp func, ModuleOp mod, MLIRContext *ctx); @@ -978,185 +961,114 @@ static std::string buildContextAttrsJson(const SpecKey &key) { } // ============================================================================ -// Invoke the configured one-shot helper and return its stdout. +// Materialize PTODSL in the host Python interpreter and import its functions. +// The service borrows the source module only for the synchronous callback; +// this pass clones the required functions into the caller module there. // ============================================================================ -std::optional -ExpandState::invokeTileLibHelper(const SpecKey &key, - StringRef candidateId) { - auto pythonPath = pto::resolvePythonExecutable(pythonExe); - if (!pythonPath) { - llvm::errs() << "ExpandTileOp: cannot find '" << pythonExe << "'\n"; - return std::nullopt; - } - - std::string operandSpecsJson = buildOperandSpecsJson(key); - std::string contextAttrsJson = buildContextAttrsJson(key); - if (key.targetArch.empty()) { - llvm::errs() << "ExpandTileOp: missing pto.target_arch module attribute\n"; - return std::nullopt; - } - - SmallString<128> tmpPath; - int tmpFD; - if (auto ec = llvm::sys::fs::createTemporaryFile("tilelib_helper", "out", - tmpFD, tmpPath)) { - llvm::errs() << "ExpandTileOp: cannot create temp file: " - << ec.message() << "\n"; - return std::nullopt; - } - ::close(tmpFD); - - std::string opName = "pto." + key.opName; - SmallVector args = { - *pythonPath, "-m", daemonHelperModule, - "--socket", daemonSocketPath, - "--target", key.targetArch, - "--op", opName, - "--operand-specs", operandSpecsJson, - }; - if (!key.contextAttrs.empty()) { - args.push_back("--context-attrs"); - args.push_back(contextAttrsJson); - } - if (!candidateId.empty()) { - args.push_back("--candidate-id"); - args.push_back(candidateId); - } +func::FuncOp ExpandState::invokeInProcessTileLib(const SpecKey &key, + StringRef candidateId, + ModuleOp mod, + MLIRContext *ctx) { + if (!tileLibService) + return nullptr; - std::optional redirects[] = {std::nullopt, StringRef(tmpPath), - std::nullopt}; - - SmallVector envp; - std::string pythonPathEnv; - std::vector envStorage; - bool hasPythonPath = !tileLibPkgPath.empty(); - if (hasPythonPath) { - const char *existingPath = ::getenv("PYTHONPATH"); - pythonPathEnv = "PYTHONPATH=" + tileLibPkgPath; - if (existingPath && existingPath[0] != '\0') { - pythonPathEnv += ":"; - pythonPathEnv += existingPath; - } - for (char **e = environ; *e; ++e) { - StringRef entry(*e); - if (entry.starts_with("PYTHONPATH=")) - continue; - envStorage.push_back(std::string(entry)); + pto::TileLibMaterializationRequest request; + request.target = key.targetArch; + request.op = "pto." + key.opName; + request.operandSpecsJson = buildOperandSpecsJson(key); + request.contextAttrsJson = buildContextAttrsJson(key); + request.candidateId = candidateId.str(); + + func::FuncOp importedEntry; + LogicalResult materializationResult = tileLibService->materialize( + request, *ctx, [&](ModuleOp sourceModule, StringRef entrySymbol) { + if (!sourceModule || sourceModule.getContext() != ctx) { + llvm::errs() << "ExpandTileOp: in-process PTODSL returned a module from " + "a different MLIRContext\n"; + return failure(); } - envStorage.push_back(pythonPathEnv); - for (auto &s : envStorage) - envp.push_back(s); - } - std::string errMsg; - int rc = llvm::sys::ExecuteAndWait( - *pythonPath, args, - hasPythonPath ? std::optional>(envp) : std::nullopt, - redirects, /*secondsToWait=*/30, /*memoryLimit=*/0, &errMsg); - - if (rc != 0) { - llvm::errs() << "ExpandTileOp: daemon helper instantiate failed (rc=" - << rc - << "): " << errMsg << "\n"; - llvm::sys::fs::remove(tmpPath); - return std::nullopt; - } + auto sourceEntry = sourceModule.lookupSymbol(entrySymbol); + if (!sourceEntry) { + llvm::errs() << "ExpandTileOp: in-process PTODSL entry symbol @" + << entrySymbol << " was not found\n"; + return failure(); + } - auto bufOrErr = llvm::MemoryBuffer::getFile(tmpPath); - llvm::sys::fs::remove(tmpPath); - if (!bufOrErr) { - llvm::errs() << "ExpandTileOp: cannot read daemon output\n"; - return std::nullopt; - } - std::string output = (*bufOrErr)->getBuffer().str(); - if (output.empty()) { - llvm::errs() << "ExpandTileOp: empty daemon output\n"; - return std::nullopt; - } - return output; -} + SmallVector sourceFuncs; + for (func::FuncOp fn : sourceModule.getOps()) + sourceFuncs.push_back(fn); + if (sourceFuncs.empty()) { + llvm::errs() << "ExpandTileOp: in-process PTODSL returned no func.func\n"; + return failure(); + } -// ============================================================================ -// Invoke the daemon RPC to generate a specialized template function. -// ============================================================================ -func::FuncOp ExpandState::invokeTileLibDaemon(const SpecKey &key, - StringRef candidateId, - ModuleOp mod, - MLIRContext *ctx) { - auto mlirText = invokeTileLibHelper(key, candidateId); - if (!mlirText) - return nullptr; + std::string uniqueName = buildUniqueFunctionBaseName(key); + if (!candidateId.empty()) + uniqueName += "__" + candidateId.str(); - // Parse the rendered MLIR. - auto parsedMod = parseSourceString(*mlirText, ctx); - if (!parsedMod) { - llvm::errs() << "ExpandTileOp: failed to parse daemon output\n"; - return nullptr; - } + SymbolTable targetSymTable(mod); + if (auto existingFunc = targetSymTable.lookup(uniqueName)) { + importedEntry = cast(existingFunc); + return success(); + } - // 9. Clone the generated function set into the target module. - auto parsedFuncs = parsedMod->getOps(); - if (parsedFuncs.empty()) { - llvm::errs() << "ExpandTileOp: no func.func in daemon output\n"; - return nullptr; - } + llvm::StringMap plannedSymbols; + for (func::FuncOp fn : sourceFuncs) { + std::string newName = fn == sourceEntry + ? uniqueName + : uniqueName + "__" + std::string(fn.getSymName()); + if (targetSymTable.lookup(newName)) { + llvm::errs() << "ExpandTileOp: imported PTODSL symbol collision at @" + << newName << "\n"; + return failure(); + } + plannedSymbols[fn.getSymName()] = std::move(newName); + } - // Create builder and set insertion point to insert functions into module - OpBuilder builder(ctx); - builder.setInsertionPointToEnd(mod.getBody()); - - llvm::StringMap renamedSymbols; - SmallVector clonedFuncs; - - std::string uniqueName = buildUniqueFunctionBaseName(key); - if (!candidateId.empty()) - uniqueName += "__" + candidateId.str(); - SymbolTable targetSymTable(mod); - if (auto existingFunc = targetSymTable.lookup(uniqueName)) - return cast(existingFunc); - - for (auto [index, fn] : llvm::enumerate(parsedFuncs)) { - // Use builder.clone() to insert into module body - IRMapping mapping; - auto cloned = cast(builder.clone(*fn, mapping)); - std::string newName; - if (index == 0) { - newName = uniqueName; - } else { - newName = uniqueName + "__" + std::string(fn.getSymName()); + OpBuilder builder(ctx); + builder.setInsertionPointToEnd(mod.getBody()); + SmallVector clonedFuncs; + for (func::FuncOp fn : sourceFuncs) { + IRMapping mapping; + auto cloned = cast(builder.clone(*fn, mapping)); + cloned.setName(plannedSymbols.lookup(fn.getSymName())); + cloned.setVisibility(SymbolTable::Visibility::Private); + clonedFuncs.push_back(cloned); } - renamedSymbols[fn.getSymName()] = newName; - cloned.setName(newName); - - // Set visibility to Private for template functions (required for inline pass) - cloned.setVisibility(SymbolTable::Visibility::Private); - - clonedFuncs.push_back(cloned); - } - for (func::FuncOp fn : clonedFuncs) { - fn.walk([&](func::CallOp call) { - StringRef callee = call.getCallee(); - if (callee.empty()) - return; - auto renameIt = renamedSymbols.find(callee); - if (renameIt == renamedSymbols.end()) - return; - call.setCallee(renameIt->second); - }); - } + for (func::FuncOp fn : clonedFuncs) { + for (const auto &renamed : plannedSymbols) { + if (failed(SymbolTable::replaceAllSymbolUses( + StringAttr::get(ctx, renamed.getKey()), + StringAttr::get(ctx, renamed.getValue()), fn))) { + llvm::errs() << "ExpandTileOp: failed to rewrite imported symbol @" + << renamed.getKey() << " in @" << fn.getSymName() + << "\n"; + for (func::FuncOp imported : clonedFuncs) + imported.erase(); + return failure(); + } + } + } - auto cloned = clonedFuncs.front(); - if (!cloned->hasAttr("pto.tilelang.instance")) { - llvm::errs() << "ExpandTileOp: warning: daemon output function @" - << cloned.getSymName() - << " missing pto.tilelang.instance attribute\n"; + importedEntry = mod.lookupSymbol(uniqueName); + if (!importedEntry) { + llvm::errs() << "ExpandTileOp: failed to import PTODSL entry @" + << entrySymbol << "\n"; + return failure(); + } + if (!importedEntry->hasAttr("pto.tilelang.instance")) + llvm::errs() << "ExpandTileOp: warning: in-process PTODSL entry @" + << importedEntry.getSymName() + << " missing pto.tilelang.instance attribute\n"; + return success(); + }); + if (failed(materializationResult)) { + llvm::errs() << "ExpandTileOp: in-process PTODSL materialization failed\n"; + return nullptr; } - - // Keep the parsed module alive. - parsedModules.push_back(std::move(parsedMod)); - - return cloned; + return importedEntry; } // ============================================================================ @@ -1165,8 +1077,9 @@ func::FuncOp ExpandState::invokeTileLibDaemon(const SpecKey &key, func::FuncOp ExpandState::invokeTileLib(const SpecKey &key, Operation *tileOp, ModuleOp mod, MLIRContext *ctx) { - if (daemonSocketPath.empty()) { - llvm::errs() << "ExpandTileOp: PTODSL backend requires its daemon\n"; + if (!tileLibService) { + tileOp->emitError( + "ExpandTileOp PTODSL backend requires an in-process service"); return nullptr; } @@ -1187,13 +1100,7 @@ func::FuncOp ExpandState::invokeTileLib(const SpecKey &key, return nullptr; } - func::FuncOp daemonResult = - invokeTileLibDaemon(key, selectedName.getValue(), mod, ctx); - if (daemonResult) - return daemonResult; - - llvm::errs() << "ExpandTileOp: PTODSL daemon RPC failed\n"; - return nullptr; + return invokeInProcessTileLib(key, selectedName.getValue(), mod, ctx); } // ============================================================================ @@ -1219,7 +1126,7 @@ LogicalResult ExpandState::expandTileOpsInFunction(func::FuncOp func, return failure(); } - // Invoke the selected TileLib backend (with daemon-side caching). + // Materialize the selected PTODSL template in-process. func::FuncOp dslFn = invokeTileLib(*specKeyOpt, op, mod, ctx); if (!dslFn) { StringRef opName = getTileOpName(op); @@ -1268,24 +1175,16 @@ void ExpandTileOpPass::runOnOperation() { if (!hasExpandableOps) return; - if (tileLibBackend != "ptodsl") { - mod.emitError("ExpandTileOp received unsupported tile-lib-backend '" + - std::string(tileLibBackend) + "'"); - signalPassFailure(); - return; - } - - if (daemonSocketPath.empty()) { - mod.emitError("ExpandTileOp requires a running PTODSL TileLib daemon"); + std::shared_ptr tileLibService = + pto::TileLibRuntime::getService(); + if (!tileLibService) { + mod.emitError("ExpandTileOp requires an initialized PTODSL runtime"); signalPassFailure(); return; } ExpandState state; - state.tileLibPkgPath = std::string(tileLibPkgPath); - state.daemonHelperModule = std::string(daemonHelperModule); - state.pythonExe = std::string(pythonExe); - state.daemonSocketPath = std::string(daemonSocketPath); + state.tileLibService = tileLibService; for (auto func : mod.getOps()) { if (func.isExternal()) @@ -1304,10 +1203,5 @@ std::unique_ptr createExpandTileOpPass() { return std::make_unique(); } -std::unique_ptr -createExpandTileOpPass(const ExpandTileOpOptions &options) { - return std::make_unique(options); -} - } // namespace pto } // namespace mlir diff --git a/lib/PTO/Transforms/InsertTemplateAttributes.cpp b/lib/PTO/Transforms/InsertTemplateAttributes.cpp index 7c6f51f808..471a015051 100644 --- a/lib/PTO/Transforms/InsertTemplateAttributes.cpp +++ b/lib/PTO/Transforms/InsertTemplateAttributes.cpp @@ -8,7 +8,6 @@ #include "PTO/IR/PTO.h" #include "PTO/IR/PTOTypeUtils.h" -#include "PTO/Support/PythonExecutable.h" #include "PTO/Transforms/Passes.h" #include "PTO/Transforms/TileOpExpansionUtils.h" @@ -21,27 +20,17 @@ #include "mlir/Pass/Pass.h" #include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SmallString.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringExtras.h" #include "llvm/ADT/StringRef.h" #include "llvm/Support/Error.h" -#include "llvm/Support/FileSystem.h" #include "llvm/Support/JSON.h" -#include "llvm/Support/MemoryBuffer.h" -#include "llvm/Support/Program.h" #include "llvm/Support/raw_ostream.h" -#include #include #include -#include #include -extern "C" { -extern char **environ; -} - using namespace mlir; namespace mlir { @@ -764,121 +753,6 @@ getTargetArch(Operation *operation) { return std::nullopt; } -static std::optional -invokeMetadataHelper(Operation *operation, StringRef pythonExe, - StringRef daemonSocketPath, StringRef tileLibPkgPath, - StringRef daemonHelperModule) { - auto pythonPath = pto::resolvePythonExecutable(pythonExe); - if (!pythonPath) { - operation->emitError("InsertTemplateAttributes cannot find Python '") - << pythonExe << "'"; - return std::nullopt; - } - - auto target = getTargetArch(operation); - auto operandSpecs = buildOperandSpecsJson(operation); - if (!target || !operandSpecs) - return std::nullopt; - std::string contextAttrs = buildContextAttrsJson(operation); - - llvm::SmallString<128> outputPath; - int outputFd; - if (auto error = llvm::sys::fs::createTemporaryFile( - "tilelib_metadata", "json", outputFd, outputPath)) { - operation->emitError("InsertTemplateAttributes cannot create temporary " - "metadata output: ") - << error.message(); - return std::nullopt; - } - ::close(outputFd); - - llvm::SmallString<128> errorPath; - int errorFd; - if (auto error = llvm::sys::fs::createTemporaryFile( - "tilelib_metadata", "err", errorFd, errorPath)) { - llvm::sys::fs::remove(outputPath); - operation->emitError("InsertTemplateAttributes cannot create temporary " - "metadata error output: ") - << error.message(); - return std::nullopt; - } - ::close(errorFd); - - std::string opName = operation->getName().getStringRef().str(); - SmallVector args = { - *pythonPath, "-m", daemonHelperModule, - "--method", "get_metadata", "--socket", - daemonSocketPath, "--target", *target, - "--op", opName, "--operand-specs", - *operandSpecs, - }; - if (contextAttrs != "{}") { - args.push_back("--context-attrs"); - args.push_back(contextAttrs); - } - - std::optional redirects[] = { - std::nullopt, - StringRef(outputPath), - StringRef(errorPath), - }; - - SmallVector environment; - std::string pythonPathEnvironment; - std::vector environmentStorage; - bool hasPythonPath = !tileLibPkgPath.empty(); - if (hasPythonPath) { - const char *existingPath = ::getenv("PYTHONPATH"); - pythonPathEnvironment = "PYTHONPATH=" + tileLibPkgPath.str(); - if (existingPath && existingPath[0] != '\0') - pythonPathEnvironment += ":" + std::string(existingPath); - - for (char **entry = environ; *entry; ++entry) { - StringRef value(*entry); - if (!value.starts_with("PYTHONPATH=")) - environmentStorage.push_back(value.str()); - } - environmentStorage.push_back(pythonPathEnvironment); - for (std::string &value : environmentStorage) - environment.push_back(value); - } - - std::string errorMessage; - int result = llvm::sys::ExecuteAndWait( - *pythonPath, args, - hasPythonPath - ? std::optional>(environment) - : std::nullopt, - redirects, /*secondsToWait=*/30, /*memoryLimit=*/0, &errorMessage); - if (result != 0) { - auto errorOutput = llvm::MemoryBuffer::getFile(errorPath); - llvm::sys::fs::remove(outputPath); - llvm::sys::fs::remove(errorPath); - - std::string detail; - if (errorOutput) - detail = errorOutput.get()->getBuffer().trim().str(); - if (detail.empty()) - detail = errorMessage; - if (detail.empty()) - detail = "helper exited with status " + std::to_string(result); - - operation->emitError("InsertTemplateAttributes metadata RPC failed: ") - << detail; - return std::nullopt; - } - - auto output = llvm::MemoryBuffer::getFile(outputPath); - llvm::sys::fs::remove(outputPath); - llvm::sys::fs::remove(errorPath); - if (!output) { - operation->emitError( - "InsertTemplateAttributes cannot read metadata output"); - return std::nullopt; - } - return (*output)->getBuffer().str(); -} - static FailureOr parseCandidateAttributes(Operation *operation, StringRef metadataJson) { auto parsed = llvm::json::parse(metadataJson); @@ -988,18 +862,29 @@ struct InsertTemplateAttributesPass }); if (tileOperations.empty()) return; - if (daemonSocketPath.empty()) { + std::shared_ptr tileLibService = + pto::TileLibRuntime::getService(); + if (!tileLibService) { module.emitError( - "InsertTemplateAttributes requires a PTODSL daemon socket"); + "InsertTemplateAttributes requires an initialized PTODSL runtime"); return signalPassFailure(); } for (Operation *operation : tileOperations) { - auto metadata = invokeMetadataHelper( - operation, pythonExe, daemonSocketPath, tileLibPkgPath, - daemonHelperModule); - if (!metadata) + auto target = getTargetArch(operation); + auto operandSpecs = buildOperandSpecsJson(operation); + if (!target || !operandSpecs) + return signalPassFailure(); + pto::TileLibMaterializationRequest request; + request.target = std::move(*target); + request.op = operation->getName().getStringRef().str(); + request.operandSpecsJson = std::move(*operandSpecs); + request.contextAttrsJson = buildContextAttrsJson(operation); + FailureOr metadata = tileLibService->getMetadata(request); + if (failed(metadata)) { + operation->emitError("in-process PTODSL metadata query failed"); return signalPassFailure(); + } auto candidates = parseCandidateAttributes(operation, *metadata); if (failed(candidates)) @@ -1018,10 +903,5 @@ std::unique_ptr createInsertTemplateAttributesPass() { return std::make_unique(); } -std::unique_ptr createInsertTemplateAttributesPass( - const InsertTemplateAttributesOptions &options) { - return std::make_unique(options); -} - } // namespace pto } // namespace mlir diff --git a/lib/PTO/Transforms/TileLibService.cpp b/lib/PTO/Transforms/TileLibService.cpp new file mode 100644 index 0000000000..3efc648381 --- /dev/null +++ b/lib/PTO/Transforms/TileLibService.cpp @@ -0,0 +1,43 @@ +// 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 "PTO/Transforms/TileLibService.h" + +#include + +namespace { + +std::mutex &getRuntimeMutex() { + static std::mutex mutex; + return mutex; +} + +std::shared_ptr &getRuntimeService() { + static std::shared_ptr service; + return service; +} + +} // namespace + +void mlir::pto::TileLibRuntime::install( + std::shared_ptr service) { + std::lock_guard lock(getRuntimeMutex()); + getRuntimeService() = std::move(service); +} + +void mlir::pto::TileLibRuntime::uninstall(TileLibService *service) { + std::lock_guard lock(getRuntimeMutex()); + if (getRuntimeService().get() == service) + getRuntimeService().reset(); +} + +std::shared_ptr +mlir::pto::TileLibRuntime::getService() { + std::lock_guard lock(getRuntimeMutex()); + return getRuntimeService(); +} diff --git a/ptodsl/README.md b/ptodsl/README.md index f4c8569cff..5e77b2a8ab 100644 --- a/ptodsl/README.md +++ b/ptodsl/README.md @@ -86,18 +86,17 @@ LLVM_BUILD_DIR=/path/to/llvm/build ./quick_install.sh ## PTODSL TileLib backend -PTOAS uses the PTODSL TileLib daemon by default for VPTO tile-op expansion: +PTOAS uses its in-process PTODSL TileLib service for VPTO tile-op expansion: ```bash ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto \ input.pto -o - ``` -Wheel and CMake-tree launchers pass the Python root containing their own -installed or staged `ptodsl` package to the native driver. Use -`--ptodsl-pkg-path=/path/to/package/root` for an explicit command-line -override. PTODSL daemon failures are reported as errors and never fall back to -the TileLang implementation. +Wheel and CMake-tree launchers add their packaged `TileOps` resource directory +to the host interpreter while the native compiler runs. Template discovery, +metadata queries, and MLIR materialization therefore stay in one process and +do not require a Python executable, daemon socket, or package-path CLI option. `InsertTemplateAttributes` queries legal-candidate metadata before fusion and stores an ordered `candidates` array containing only `id`, `name`, diff --git a/ptodsl/docs/developer_guide/tilelib-debugging-playbook.md b/ptodsl/docs/developer_guide/tilelib-debugging-playbook.md index d5c4c015f8..8b5d7775c8 100644 --- a/ptodsl/docs/developer_guide/tilelib-debugging-playbook.md +++ b/ptodsl/docs/developer_guide/tilelib-debugging-playbook.md @@ -1,7 +1,7 @@ # PTODSL TileLib Debugging Playbook -This playbook is for diagnosing PTODSL TileLib failures while migrating from -TileLangDSL. It converts the migration scratch notes into a reusable workflow. +This playbook is for diagnosing PTODSL TileLib failures. It converts the +original migration scratch notes into a reusable workflow. The main rule: classify the failure before editing templates. A build failure, a candidate-selection failure, a tracing failure, and a wrong-output failure @@ -24,7 +24,6 @@ ninja -C build-llvm21 PTODSLPackage Run one smoke ST: ```bash -PTOAS_TILE_LIB_BACKEND=ptodsl \ python3 test/tilelang_st/script/run_all_st.py \ -r sim -v a5 \ -p build-llvm21/tools/ptoas/ptoas \ @@ -34,7 +33,6 @@ python3 test/tilelang_st/script/run_all_st.py \ Run one non-smoke ST: ```bash -PTOAS_TILE_LIB_BACKEND=ptodsl \ python3 test/tilelang_st/script/run_all_st.py \ -r sim -v a5 \ -p build-llvm21/tools/ptoas/ptoas \ @@ -44,7 +42,6 @@ python3 test/tilelang_st/script/run_all_st.py \ Run one named ST case when supported by `run_st.py`: ```bash -PTOAS_TILE_LIB_BACKEND=ptodsl \ python3 test/tilelang_st/script/run_st.py \ -r sim -v a5 \ -p build-llvm21/tools/ptoas/ptoas \ @@ -64,8 +61,7 @@ python3 test/tilelang_st/script/run_st.py \ | isolated case passes but full test fails | helper specialization cache or stale generated package | Do not assume every ST failure means the testcase is wrong. First check whether -the same case works with TileLangDSL and whether the PTODSL lowering has enough -metadata to reproduce the TileLangDSL behavior. +the PTODSL lowering has enough metadata to represent the testcase behavior. ## Candidate Selection Failures @@ -94,15 +90,15 @@ If legality appears correct but `ExpandTileOp` cannot expand: non-empty `candidates` attr. 2. Dump after passes that rewrite view/tile operands and confirm the attr is still attached. -3. Confirm candidate 0 has a `name` and that the daemon can render that - candidate directly. +3. Confirm candidate 0 has a `name` and that the in-process TileLib service can + materialize that candidate directly. Useful compiler-only command: ```bash build-llvm21/tools/ptoas/ptoas \ --pto-arch=a5 --pto-backend=vpto --emit-vpto \ - --tile-lib-backend=ptodsl --enable-insert-sync \ + --enable-insert-sync \ --mlir-print-ir-after=pto-expand-tile-op \ --mlir-print-ir-tree-dir=/tmp/_after_expand_ptodsl \ test/tilelang_st/npu/a5/src/st/testcase//.pto \ @@ -133,7 +129,7 @@ Compiler-only dump: ```bash build-llvm21/tools/ptoas/ptoas \ --pto-arch=a5 --pto-backend=vpto --emit-vpto \ - --tile-lib-backend=ptodsl --enable-insert-sync \ + --enable-insert-sync \ test/tilelang_st/npu/a5/src/st/testcase//.pto \ -o /tmp/_ptodsl.vpto ``` diff --git a/ptodsl/docs/developer_guide/tilelib-template-authoring.md b/ptodsl/docs/developer_guide/tilelib-template-authoring.md index df05cefd6c..2eca3b56e0 100644 --- a/ptodsl/docs/developer_guide/tilelib-template-authoring.md +++ b/ptodsl/docs/developer_guide/tilelib-template-authoring.md @@ -33,7 +33,7 @@ def template_tadd(src0, src1, dst): ... ``` -The function parameter order is the operand binding contract. The daemon binds +The function parameter order is the operand binding contract. The TileLib runtime binds MLIR operands positionally to these parameter names before evaluating constraints or rendering. If a TileLangDSL template had multiple callable forms, either match the ST operand order exactly or register separate PTODSL diff --git a/ptodsl/ptoas/_cli.py b/ptodsl/ptoas/_cli.py index a12158ad7d..90e5b5ecc0 100644 --- a/ptodsl/ptoas/_cli.py +++ b/ptodsl/ptoas/_cli.py @@ -31,13 +31,12 @@ def _load_native_module(): return _core -def _resolve_runtime_paths(native_module) -> tuple[Path, Path]: +def _resolve_tileops_dir(native_module) -> Path: module_file = getattr(native_module, "__file__", None) if not module_file: raise SystemExit("ptoas._core does not expose a module file") package_root = Path(module_file).resolve().parent - python_root = package_root.parent runtime_root = package_root / "_runtime" tileops_dir = runtime_root / "share" / "ptoas" / "TileOps" if not tileops_dir.is_dir(): @@ -45,32 +44,27 @@ def _resolve_runtime_paths(native_module) -> tuple[Path, Path]: "unable to locate packaged PTOAS TileOps resources: expected " f"{tileops_dir}" ) - return python_root, tileops_dir.resolve() - - -def _has_cli_option(arguments: Sequence[str], option: str) -> bool: - option_with_value = f"{option}=" - return any( - argument == option or argument.startswith(option_with_value) - for argument in arguments - ) + return tileops_dir.resolve() def launch(user_args: Sequence[str], *, wrapper: Path | None = None) -> int: native_module = _load_native_module() - python_root, tileops_dir = _resolve_runtime_paths(native_module) + tileops_dir = _resolve_tileops_dir(native_module) wrapper = wrapper.resolve() if wrapper is not None else _resolve_wrapper_path() os.environ["PTOAS_BIN"] = str(wrapper) - os.environ["PTOAS_PYTHON_EXE"] = sys.executable argv = [str(wrapper)] - if not _has_cli_option(user_args, "--ptodsl-pkg-path"): - argv.extend(["--ptodsl-pkg-path", str(python_root)]) - if not _has_cli_option(user_args, "--tileops-pkg-path"): - argv.extend(["--tileops-pkg-path", str(tileops_dir.parent)]) argv.extend(user_args) - return int(native_module.main(argv)) + tileops_python_root = str(tileops_dir.parent) + inserted_tileops_root = tileops_python_root not in sys.path + if inserted_tileops_root: + sys.path.insert(0, tileops_python_root) + try: + return int(native_module.main(argv)) + finally: + if inserted_tileops_root: + sys.path.remove(tileops_python_root) def main() -> int: diff --git a/ptodsl/ptodsl/_surface_values.py b/ptodsl/ptodsl/_surface_values.py index 0147e9bda7..375806b615 100644 --- a/ptodsl/ptodsl/_surface_values.py +++ b/ptodsl/ptodsl/_surface_values.py @@ -520,34 +520,40 @@ class _TileValidShapeView: def __init__(self, tile: "TileValue"): self._tile = tile - self._cache: dict[int, object] = {} def __getitem__(self, index: int): - logical_rank = len(self._tile.shape) if self._tile.shape is not None else 2 + return self._tile._get_valid_shape_dim(index) + + +class TileValue(_SurfaceValue, Tile): + """Author-facing tile handle with surface-style accessors.""" + + def _get_valid_shape_dim(self, index: int): + logical_rank = len(self.shape) if self.shape is not None else 2 allowed = {0} if logical_rank == 1 else {0, 1} if index not in allowed: if logical_rank == 1: raise IndexError("PTODSL rank-1 tile.valid_shape currently supports only index 0") raise IndexError("PTODSL tile.valid_shape currently supports indices 0 and 1") - cached = self._cache.get(index) + cached = self._valid_shape_cache.get(index) if cached is not None: return cached - if self._tile.static_valid_shape is not None: - dim = self._tile.static_valid_shape[index] + if self.static_valid_shape is not None: + dim = self.static_valid_shape[index] if dim is not None: value = _index_const(dim) if _is_python_index_literal(dim) else unwrap_surface_value(dim) value = wrap_surface_value(value) - self._cache[index] = value + self._valid_shape_cache[index] = value return value try: if logical_rank == 1: - value = wrap_surface_value(_pto.TileValidColsOp(self._tile.value).result) + value = wrap_surface_value(_pto.TileValidColsOp(self.value).result) elif index == 0: - value = wrap_surface_value(_pto.TileValidRowsOp(self._tile.value).result) + value = wrap_surface_value(_pto.TileValidRowsOp(self.value).result) else: - value = wrap_surface_value(_pto.TileValidColsOp(self._tile.value).result) + value = wrap_surface_value(_pto.TileValidColsOp(self.value).result) except Exception: - static_dim = _fallback_static_valid_dim(self._tile.type, index) + static_dim = _fallback_static_valid_dim(self.type, index) if static_dim is None: raise RuntimeError( "tile.valid_shape could not be lowered because the current " @@ -555,13 +561,9 @@ def __getitem__(self, index: int): "the tile type does not carry a recoverable static bound" ) from None value = wrap_surface_value(_index_const(static_dim)) - self._cache[index] = value + self._valid_shape_cache[index] = value return value - -class TileValue(_SurfaceValue, Tile): - """Author-facing tile handle with surface-style accessors.""" - def __init__( self, value, @@ -591,11 +593,11 @@ def __init__( self.static_valid_shape = tuple(valid_shape) if valid_shape is not None else ( parsed["valid_dims"] if parsed is not None else None ) - self._valid_shape = _TileValidShapeView(self) + self._valid_shape_cache: dict[int, object] = {} @property def valid_shape(self): - return self._valid_shape + return _TileValidShapeView(self) @valid_shape.setter def valid_shape(self, dims): @@ -603,7 +605,7 @@ def valid_shape(self, dims): set_tile_valid_shape(self, dims) self.static_valid_shape = tuple(dims) - self._valid_shape._cache.clear() + self._valid_shape_cache.clear() @property def surface_metadata(self): diff --git a/ptodsl/ptodsl/tilelib/_compiler_runtime.py b/ptodsl/ptodsl/tilelib/_compiler_runtime.py new file mode 100644 index 0000000000..3fbcaa29f2 --- /dev/null +++ b/ptodsl/ptodsl/tilelib/_compiler_runtime.py @@ -0,0 +1,73 @@ +# 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. + +"""In-process PTODSL TileLib materialization entry point.""" + +from __future__ import annotations + +import json + +from ._render_runtime import _TemplateTrace +from ._selection import _select_descriptor_and_specs, metadata_request + + +def metadata( + target: str, + op: str, + operand_specs_json: str, + context_attrs_json: str, +) -> str: + """Return candidate metadata JSON without any daemon transport.""" + try: + operand_specs = json.loads(operand_specs_json) + context_attrs = json.loads(context_attrs_json or "{}") + except json.JSONDecodeError as exc: + raise ValueError(f"invalid TileLib metadata request: {exc}") from exc + return json.dumps( + metadata_request(target, op, operand_specs, context_attrs), + separators=(",", ":"), + sort_keys=True, + ) + + +def materialize( + target: str, + op: str, + operand_specs_json: str, + context_attrs_json: str, + candidate_id: str | None, + context, +): + """Return ``(source_module, entry_symbol)`` in *context*. + + JSON is retained only for the compact pure-data specialization request. No + MLIR text is produced or parsed by this path. + """ + try: + operand_specs = json.loads(operand_specs_json) + context_attrs = json.loads(context_attrs_json or "{}") + except json.JSONDecodeError as exc: + raise ValueError(f"invalid TileLib materialization request: {exc}") from exc + + descriptor, tile_specs = _select_descriptor_and_specs( + target, + op, + operand_specs, + context_attrs, + candidate_id or None, + ) + module = _TemplateTrace( + descriptor, + tile_specs, + context_attrs=context_attrs, + ).build_module_in_context(context) + module.operation.verify() + return module, descriptor.name + + +__all__ = ["materialize", "metadata"] diff --git a/ptodsl/ptodsl/tilelib/_render_runtime.py b/ptodsl/ptodsl/tilelib/_render_runtime.py index d925c76ce0..bd23dcdab0 100644 --- a/ptodsl/ptodsl/tilelib/_render_runtime.py +++ b/ptodsl/ptodsl/tilelib/_render_runtime.py @@ -69,7 +69,7 @@ def __init__(self, value, spec: TileSpec): ) # Force the dynamic valid-shape ops to match the tilelang render. self.static_valid_shape = None - self._valid_shape._cache.clear() + self._valid_shape_cache.clear() self._template_static_valid_shape = tuple(spec.valid_shape or spec.shape) self._template_config = _TemplateTileConfig( b_layout=spec.b_layout, @@ -176,9 +176,15 @@ def trace_entry(self, *args): rewritten(*args) # Custom golden-shaped container: single module(target_arch) + func(instance, kernel_kind). - def build_module(self): - ctx = make_context() - with ctx, Location.unknown(): + def build_standalone_module(self): + """Build a module in a fresh context for standalone PTODSL use.""" + return self.build_module_in_context(make_context()) + + def build_module_in_context(self, context): + """Build a compiler-owned source module in the explicit context.""" + if context is None: + raise TypeError("compiler materialization requires an explicit context") + with context, Location.unknown(): arg_types = list(self.compute_argument_types()) module, ir_fn = self._create_instance_module(arg_types) session = self.create_session(module, ir_fn) @@ -192,6 +198,10 @@ def build_module(self): self.finalize_session(session) session.validate_final_state() self.verify_module(module) + if module.context is not context: + raise RuntimeError( + "TileLib materialization returned a module from a different context" + ) return module def _create_instance_module(self, arg_types): diff --git a/ptodsl/ptodsl/tilelib/serving/daemon.py b/ptodsl/ptodsl/tilelib/_selection.py similarity index 58% rename from ptodsl/ptodsl/tilelib/serving/daemon.py rename to ptodsl/ptodsl/tilelib/_selection.py index f52c4c6dcf..d34ebe51f7 100644 --- a/ptodsl/ptodsl/tilelib/serving/daemon.py +++ b/ptodsl/ptodsl/tilelib/_selection.py @@ -5,43 +5,18 @@ # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. -"""PTODSL TileLib daemon for the ExpandTileOp Unix-socket RPC contract. - -The daemon owns template discovery, selection, specialization, rendering, and an -in-memory instance cache. PTODSL templates are loaded from the Python package, so -the daemon does not scan or depend on an external template directory. - -Run it with: - - python3 -m ptodsl.tilelib.serving.daemon --socket -""" +"""PTODSL TileLib candidate discovery, validation, and selection.""" from __future__ import annotations -import argparse -import json -import os -import signal -import socketserver -import threading - -from .. import constraints as _constraints -from .. import registry as _registry -from ..metadata import ScalarSpec, ScalarType, TileSpec, VectorSpec, ViewSpec +from . import constraints as _constraints +from . import registry as _registry +from .metadata import ScalarSpec, ScalarType, TileSpec, VectorSpec, ViewSpec from TileOps import load_template -from .wire import recv_message, send_message - - -def _remove_socket_path(socket_path: str) -> None: - """Remove an existing socket entry, including a broken symlink.""" - try: - os.unlink(socket_path) - except FileNotFoundError: - pass def _build_tile_specs(descriptor, operand_specs: list) -> dict: - """Map positional daemon operands onto a template's parameter names.""" + """Map positional compiler operands onto a template's parameter names.""" if not isinstance(operand_specs, list): raise TypeError("operand_specs must be a list") if len(operand_specs) != len(descriptor.param_names): @@ -101,7 +76,7 @@ def _build_tile_specs(descriptor, operand_specs: list) -> dict: if kind != "tile": raise NotImplementedError( - "PTODSL TileLib daemon currently supports tile, scalar, view, " + "PTODSL TileLib currently supports tile, scalar, view, " f"and vector operands; " f"operand {index} ({name!r}) has kind {kind!r}" ) @@ -296,183 +271,4 @@ def metadata_request( } -def render_request( - target: str, - op: str, - operand_specs: list, - context_attrs: dict | None = None, - candidate_id: str | None = None, -) -> str: - """Select and render one PTODSL template as MLIR text.""" - descriptor, tile_specs = _select_descriptor_and_specs( - target, - op, - operand_specs, - context_attrs, - candidate_id, - ) - return descriptor.specialize( - context_attrs=context_attrs or {}, - **tile_specs, - ).mlir_text() - - -class TileLibDaemonServer(socketserver.UnixStreamServer): - """Sequential Unix-socket RPC server with an in-memory render cache.""" - - def __init__(self, socket_path: str, max_entries: int = 1000): - if max_entries <= 0: - raise ValueError("max_entries must be greater than zero") - super().__init__(socket_path, _Handler) - os.chmod(socket_path, 0o600) - self._cache: dict[str, str] = {} - self._max_entries = max_entries - self._stats = {"hits": 0, "misses": 0, "evictions": 0} - - @property - def stats(self) -> dict: - """Return a snapshot of cache counters for diagnostics and tests.""" - return dict(self._stats) - - def dispatch(self, request: dict) -> dict: - if not isinstance(request, dict): - return {"success": False, "error": "request must be a JSON object"} - - method = request.get("method") - params = request.get("params") or {} - if not isinstance(params, dict): - return {"success": False, "error": "request params must be a JSON object"} - - try: - if method == "instantiate": - result = self._instantiate(**params) - elif method == "get_metadata": - result = self._get_metadata(**params) - elif method == "ping": - result = "pong" - elif method == "get_stats": - result = self._get_stats() - elif method == "clear": - result = self._clear() - else: - return {"success": False, "error": f"unknown method {method!r}"} - return {"success": True, "result": result} - except Exception as exc: - return { - "success": False, - "error": f"{type(exc).__name__}: {exc}", - } - - def _get_metadata(self, target, op, operand_specs, context_attrs=None): - return metadata_request(target, op, operand_specs, context_attrs) - - def _get_stats(self): - requests = self._stats["hits"] + self._stats["misses"] - total_entries = len(self._cache) - return { - **self._stats, - "entries": total_entries, - "total_entries": total_entries, - "max_entries": self._max_entries, - "hit_rate": self._stats["hits"] / requests if requests else 0.0, - } - - def _clear(self): - self._cache.clear() - return {"cleared": True} - - def _instantiate( - self, - target, - op, - operand_specs, - context_attrs=None, - candidate_id=None, - ): - key = json.dumps( - { - "target": target, - "op": op, - "operand_specs": operand_specs, - "context_attrs": context_attrs, - "candidate_id": candidate_id, - }, - sort_keys=True, - separators=(",", ":"), - ) - - cached = self._cache.get(key) - if cached is not None: - self._stats["hits"] += 1 - return cached - self._stats["misses"] += 1 - - mlir_text = render_request( - target, - op, - operand_specs, - context_attrs, - candidate_id, - ) - - if len(self._cache) >= self._max_entries: - self._cache.pop(next(iter(self._cache))) - self._stats["evictions"] += 1 - self._cache[key] = mlir_text - return mlir_text - - -class _Handler(socketserver.BaseRequestHandler): - def handle(self): - try: - request = recv_message(self.request) - except (ConnectionError, UnicodeDecodeError, ValueError): - return - send_message(self.request, self.server.dispatch(request)) - - -def _parse_args(argv): - parser = argparse.ArgumentParser(prog="ptodsl.tilelib.serving.daemon") - parser.add_argument("--socket", required=True) - parser.add_argument( - "--template-dir", - default=None, - help="accepted during migration but ignored; PTODSL templates are in-package", - ) - parser.add_argument("--max-entries", type=int, default=1000) - parser.add_argument("--verbose", action="store_true") - return parser.parse_args(argv) - - -def main(argv=None): - args = _parse_args(argv) - - _remove_socket_path(args.socket) - - server = TileLibDaemonServer(args.socket, max_entries=args.max_entries) - stop = threading.Event() - - def _request_shutdown(*_): - stop.set() - - signal.signal(signal.SIGTERM, _request_shutdown) - signal.signal(signal.SIGINT, _request_shutdown) - - thread = threading.Thread(target=server.serve_forever, daemon=True) - thread.start() - if args.verbose: - print(f"PTODSL TileLib daemon listening on {args.socket}", flush=True) - - try: - stop.wait() - finally: - server.shutdown() - server.server_close() - _remove_socket_path(args.socket) - - -if __name__ == "__main__": - main() - - -__all__ = ["TileLibDaemonServer", "main", "metadata_request", "render_request"] +__all__ = ["metadata_request"] diff --git a/ptodsl/ptodsl/tilelib/decorator.py b/ptodsl/ptodsl/tilelib/decorator.py index ba0f3db095..415346d56f 100644 --- a/ptodsl/ptodsl/tilelib/decorator.py +++ b/ptodsl/ptodsl/tilelib/decorator.py @@ -50,13 +50,12 @@ def __init__(self, descriptor: TileTemplate, tile_specs: dict, context_attrs=Non descriptor.name, module_factory=lambda: _TemplateTrace( descriptor, tile_specs, context_attrs=context_attrs - ).build_module(), + ).build_standalone_module(), ) self.descriptor = descriptor self.tile_specs = tile_specs self.context_attrs = dict(context_attrs or {}) - def tile_template(*, op, target="a5", name=None, dtypes=(), layouts=(), memory_spaces=(), constraints=(), priority=0, fusible=False, loop_depth=None, id=None, Tail=None, is_post_update=False, diff --git a/ptodsl/ptodsl/tilelib/serving/__init__.py b/ptodsl/ptodsl/tilelib/serving/__init__.py deleted file mode 100644 index 3ac3c8f6a3..0000000000 --- a/ptodsl/ptodsl/tilelib/serving/__init__.py +++ /dev/null @@ -1,33 +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. -"""Unix-socket serving layer for the PTODSL TileLib.""" - -from .client import DaemonClient, DaemonError - - -def __getattr__(name): - # Keep daemon.py unloaded when executing it with ``python -m``. - if name in {"TileLibDaemonServer", "metadata_request", "render_request"}: - from .daemon import TileLibDaemonServer, metadata_request, render_request - - exports = { - "TileLibDaemonServer": TileLibDaemonServer, - "metadata_request": metadata_request, - "render_request": render_request, - } - return exports[name] - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - - -__all__ = [ - "DaemonClient", - "DaemonError", - "TileLibDaemonServer", - "metadata_request", - "render_request", -] diff --git a/ptodsl/ptodsl/tilelib/serving/client.py b/ptodsl/ptodsl/tilelib/serving/client.py deleted file mode 100644 index 9a4db37b0e..0000000000 --- a/ptodsl/ptodsl/tilelib/serving/client.py +++ /dev/null @@ -1,77 +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. -"""Synchronous client for the PTODSL TileLib daemon.""" - -from __future__ import annotations - -import socket - -from .wire import recv_message, send_message - - -class DaemonError(Exception): - """An RPC reached the daemon but the requested operation failed.""" - - -class DaemonClient: - """Issue one daemon RPC per Unix-socket connection.""" - - def __init__(self, socket_path: str): - self.socket_path = socket_path - - def _call(self, method: str, params: dict | None = None): - with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock: - sock.connect(self.socket_path) - send_message(sock, {"method": method, "params": params or {}}) - response = recv_message(sock) - - if not response.get("success"): - raise DaemonError(response.get("error", "unknown daemon error")) - return response["result"] - - def ping(self): - return self._call("ping") - - def get_metadata(self, target, op, operand_specs, context_attrs=None): - return self._call( - "get_metadata", - { - "target": target, - "op": op, - "operand_specs": operand_specs, - "context_attrs": context_attrs or {}, - }, - ) - - def instantiate( - self, - target, - op, - operand_specs, - context_attrs=None, - candidate_id=None, - ): - return self._call( - "instantiate", - { - "target": target, - "op": op, - "operand_specs": operand_specs, - "context_attrs": context_attrs or {}, - "candidate_id": candidate_id, - }, - ) - - def get_stats(self): - return self._call("get_stats") - - def clear(self): - return self._call("clear") - - -__all__ = ["DaemonClient", "DaemonError"] diff --git a/ptodsl/ptodsl/tilelib/serving/helper.py b/ptodsl/ptodsl/tilelib/serving/helper.py deleted file mode 100644 index 5925556ef7..0000000000 --- a/ptodsl/ptodsl/tilelib/serving/helper.py +++ /dev/null @@ -1,73 +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. -"""One-shot command-line client for the ExpandTileOp daemon contract. - -Example: - - python3 -m ptodsl.tilelib.serving.helper --socket --target a5 \ - --op pto.tadd --operand-specs '[...]' -""" - -from __future__ import annotations - -import argparse -import json -import sys - -from .client import DaemonClient, DaemonError - - -def main(argv=None): - parser = argparse.ArgumentParser(prog="ptodsl.tilelib.serving.helper") - parser.add_argument("--socket", required=True) - parser.add_argument("--target", required=True) - parser.add_argument("--op", required=True) - parser.add_argument("--operand-specs", required=True) - parser.add_argument("--context-attrs", default=None) - parser.add_argument( - "--method", - choices=("instantiate", "get_metadata"), - default="instantiate", - ) - parser.add_argument("--candidate-id", default=None) - args = parser.parse_args(argv) - - try: - operand_specs = json.loads(args.operand_specs) - context_attrs = json.loads(args.context_attrs) if args.context_attrs else {} - except json.JSONDecodeError as exc: - parser.error(f"invalid JSON input: {exc}") - - try: - client = DaemonClient(args.socket) - if args.method == "get_metadata": - result = client.get_metadata( - args.target, - args.op, - operand_specs, - context_attrs, - ) - sys.stdout.write(json.dumps(result)) - return - - result = client.instantiate( - args.target, - args.op, - operand_specs, - context_attrs, - args.candidate_id, - ) - except (DaemonError, OSError) as exc: - sys.stderr.write(f"Error: daemon RPC failed: {exc}\n") - raise SystemExit(1) from exc - - sys.stdout.write(result) - - -if __name__ == "__main__": - main() diff --git a/ptodsl/ptodsl/tilelib/serving/wire.py b/ptodsl/ptodsl/tilelib/serving/wire.py deleted file mode 100644 index 8f9185137b..0000000000 --- a/ptodsl/ptodsl/tilelib/serving/wire.py +++ /dev/null @@ -1,57 +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. -"""Length-prefixed JSON framing for the TileLib daemon RPC.""" - -from __future__ import annotations - -import json - - -MAX_MESSAGE_SIZE = 64 * 1024 * 1024 - - -def recv_exactly(sock, length: int) -> bytes: - """Read exactly ``length`` bytes or fail if the peer closes early.""" - chunks = [] - remaining = length - while remaining: - chunk = sock.recv(remaining) - if not chunk: - raise ConnectionError("socket closed mid-message") - chunks.append(chunk) - remaining -= len(chunk) - return b"".join(chunks) - - -def send_message(sock, message: dict) -> None: - """Send one UTF-8 JSON message with a 4-byte big-endian length prefix.""" - payload = json.dumps(message).encode("utf-8") - if len(payload) > MAX_MESSAGE_SIZE: - raise ValueError( - f"message length {len(payload)} exceeds limit {MAX_MESSAGE_SIZE}" - ) - sock.sendall(len(payload).to_bytes(4, byteorder="big")) - sock.sendall(payload) - - -def recv_message(sock) -> dict: - """Receive one length-prefixed UTF-8 JSON message.""" - length = int.from_bytes(recv_exactly(sock, 4), byteorder="big") - if length > MAX_MESSAGE_SIZE: - raise ValueError( - f"message length {length} exceeds limit {MAX_MESSAGE_SIZE}" - ) - return json.loads(recv_exactly(sock, length).decode("utf-8")) - - -__all__ = [ - "MAX_MESSAGE_SIZE", - "recv_exactly", - "recv_message", - "send_message", -] diff --git a/ptodsl/tests/test_ptoas_cli.py b/ptodsl/tests/test_ptoas_cli.py index 51aa3c07dd..fd30b86f71 100644 --- a/ptodsl/tests/test_ptoas_cli.py +++ b/ptodsl/tests/test_ptoas_cli.py @@ -45,20 +45,13 @@ def test_launch_uses_standard_native_module_and_packaged_resources(self): self.assertEqual(exit_code, 0) native_module.main.assert_called_once_with( - [ - str(wrapper.resolve()), - "--ptodsl-pkg-path", - str(package_root.parent.resolve()), - "--tileops-pkg-path", - str(tileops_dir.parent.resolve()), - "--version", - ] + [str(wrapper.resolve()), "--version"] ) self.assertEqual(environment["PTOAS_BIN"], str(wrapper.resolve())) - self.assertEqual(environment["PTOAS_PYTHON_EXE"], _cli.sys.executable) + self.assertNotIn("PTOAS_PYTHON_EXE", environment) self.assertEqual(environment["PATH"], "/usr/bin") - def test_explicit_resource_options_are_not_overridden(self): + def test_user_arguments_are_forwarded_unchanged(self): with tempfile.TemporaryDirectory() as temp_dir: package_root = Path(temp_dir) / "install" / "ptoas" (package_root / "_runtime" / "share" / "ptoas" / "TileOps").mkdir( @@ -69,9 +62,7 @@ def test_explicit_resource_options_are_not_overridden(self): wrapper.write_text("", encoding="utf-8") native_module = self._make_native_module(package_root) arguments = [ - "--ptodsl-pkg-path=/custom/ptodsl", - "--tileops-pkg-path", - "/custom/tileops", + "--pto-arch=a5", "--version", ] @@ -92,11 +83,8 @@ def test_build_tree_uses_the_same_packaged_resource_layout(self): tileops_dir.mkdir(parents=True) native_module = self._make_native_module(package_root) - python_root, resolved_tileops = _cli._resolve_runtime_paths( - native_module - ) + resolved_tileops = _cli._resolve_tileops_dir(native_module) - self.assertEqual(python_root, package_root.parent.resolve()) self.assertEqual(resolved_tileops, tileops_dir.resolve()) def test_missing_tileops_resources_is_an_error(self): @@ -105,7 +93,7 @@ def test_missing_tileops_resources_is_an_error(self): native_module = self._make_native_module(package_root) with self.assertRaisesRegex(SystemExit, "TileOps"): - _cli._resolve_runtime_paths(native_module) + _cli._resolve_tileops_dir(native_module) if __name__ == "__main__": diff --git a/ptodsl/tests/test_ptoas_runtime.py b/ptodsl/tests/test_ptoas_runtime.py new file mode 100644 index 0000000000..2aa7db4258 --- /dev/null +++ b/ptodsl/tests/test_ptoas_runtime.py @@ -0,0 +1,51 @@ +#!/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 tempfile +import unittest +from pathlib import Path + +from ptoas import _core + + +INPUT = ( + Path(__file__).parents[2] + / "test" + / "lit" + / "vpto" + / "expand_tile_op_ptodsl_tadd.pto" +) + + +class PTOASRuntimeTest(unittest.TestCase): + def test_process_runtime_serves_consecutive_compilation_contexts(self): + self.assertTrue(INPUT.exists(), f"missing test input {INPUT}") + + with tempfile.TemporaryDirectory() as temp_dir: + for index in range(2): + output = Path(temp_dir) / f"result-{index}.mlir" + result = _core.main( + [ + "ptoas", + "--pto-arch=a5", + "--pto-backend=vpto", + "--emit-vpto", + str(INPUT), + "-o", + str(output), + ] + ) + + self.assertEqual(result, 0) + self.assertTrue(output.exists()) + self.assertIn("pto.vadd", output.read_text(encoding="utf-8")) + + +if __name__ == "__main__": + unittest.main() diff --git a/ptodsl/tests/test_tilelib_daemon.py b/ptodsl/tests/test_tilelib_daemon.py deleted file mode 100644 index 5d565d7a9c..0000000000 --- a/ptodsl/tests/test_tilelib_daemon.py +++ /dev/null @@ -1,266 +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. -"""End-to-end tests for the PTODSL TileLib daemon's Unix-socket RPC.""" - -import os -import socket -import stat -import tempfile -import threading -import unittest - -from ptodsl.tilelib.serving.client import DaemonClient, DaemonError -from ptodsl.tilelib.serving.daemon import ( - TileLibDaemonServer, - _remove_socket_path, -) -from ptodsl.tilelib.serving.wire import MAX_MESSAGE_SIZE, recv_message - - -def _tile_spec(dtype="f32", shape=(8, 64)): - return { - "kind": "tile", - "dtype": dtype, - "shape": list(shape), - "valid_shape": list(shape), - "memory_space": "ub", - "config": { - "b_layout": "row_major", - "s_layout": "none_box", - "s_fractal_size": 512, - "pad_value": "0x0", - }, - } - - -def _view_spec(dtype="f32", shape=(1, 1, 1, 8, 64), strides=(512, 512, 512, 64, 1)): - return { - "kind": "view", - "dtype": dtype, - "shape": list(shape), - "strides": list(strides), - "memory_space": "gm", - } - - -# ExpandTileOp sends tadd as ins(src0, src1), outs(dst), matching the -# template parameter order (src0, src1, dst). -TADD_OPERANDS = [_tile_spec(), _tile_spec(), _tile_spec()] -TADD = "template_tadd" - - -class TileLibDaemonTest(unittest.TestCase): - def setUp(self): - self._temporary_directory = tempfile.TemporaryDirectory() - self.socket_path = os.path.join( - self._temporary_directory.name, - "ptodsl_tilelib.sock", - ) - self.server = TileLibDaemonServer(self.socket_path) - self._thread = threading.Thread( - target=self.server.serve_forever, - daemon=True, - ) - self._thread.start() - self.client = DaemonClient(self.socket_path) - - def tearDown(self): - self.server.shutdown() - self.server.server_close() - self._thread.join() - self._temporary_directory.cleanup() - - def test_ping(self): - self.assertEqual(self.client.ping(), "pong") - - def test_socket_is_accessible_only_by_owner(self): - mode = stat.S_IMODE(os.stat(self.socket_path).st_mode) - self.assertEqual(mode, 0o600) - - def test_instantiate_named_candidate_returns_structured_mlir(self): - mlir = self.client.instantiate( - "a5", - "pto.tadd", - TADD_OPERANDS, - candidate_id=TADD, - ) - self.assertIn(f"func.func @{TADD}", mlir) - for operation in ( - "pto.tile_buf_addr", - "!pto.ptr", - "pto.vlds", - "pto.vadd", - "pto.vsts", - "pto.plt_b32", - "pto.tilelang.instance", - ): - self.assertIn(operation, mlir) - self.assertNotIn("pto.castptr", mlir) - - def test_instantiate_uses_single_tadd_candidate_without_explicit_id(self): - mlir = self.client.instantiate("a5", "pto.tadd", TADD_OPERANDS) - self.assertIn(f"func.func @{TADD}", mlir) - - def test_get_metadata_returns_legal_candidates(self): - metadata = self.client.get_metadata("a5", "pto.tadd", TADD_OPERANDS) - candidates = metadata["candidates"] - self.assertEqual( - set(candidates), - {TADD}, - ) - - selected = candidates[TADD] - self.assertEqual(selected["loop_depth"], 2) - self.assertIsNone(selected["Tail"]) - self.assertFalse(selected["has_tail"]) - self.assertFalse(selected["is_post_update"]) - self.assertEqual(selected["iteration_axis"], "none") - self.assertEqual(selected["op_engine"], "vector") - self.assertEqual(selected["op_class"], "elementwise") - self.assertEqual(selected["tags"], ["elementwise", "binary"]) - - def test_cache_stats_and_clear_are_available_over_rpc(self): - arguments = ( - "a5", - "pto.tadd", - TADD_OPERANDS, - ) - self.client.instantiate( - *arguments, - candidate_id=TADD, - ) - self.client.instantiate( - *arguments, - candidate_id=TADD, - ) - - stats = self.client.get_stats() - self.assertEqual(stats["misses"], 1) - self.assertEqual(stats["hits"], 1) - self.assertEqual(stats["entries"], 1) - - self.assertEqual(self.client.clear(), {"cleared": True}) - self.assertEqual(self.client.get_stats()["entries"], 0) - - def test_cache_key_includes_context_attributes(self): - self.client.instantiate( - "a5", - "pto.tadd", - TADD_OPERANDS, - context_attrs={"variant": 0}, - candidate_id=TADD, - ) - self.client.instantiate( - "a5", - "pto.tadd", - TADD_OPERANDS, - context_attrs={"variant": 1}, - candidate_id=TADD, - ) - self.assertEqual(self.client.get_stats()["misses"], 2) - - def test_oversized_wire_message_is_rejected_before_payload_read(self): - receiver, sender = socket.socketpair() - self.addCleanup(receiver.close) - self.addCleanup(sender.close) - sender.sendall((MAX_MESSAGE_SIZE + 1).to_bytes(4, byteorder="big")) - - with self.assertRaisesRegex(ValueError, "exceeds limit"): - recv_message(receiver) - - def test_socket_cleanup_removes_broken_symlink(self): - missing_target = os.path.join( - self._temporary_directory.name, - "missing.sock", - ) - broken_link = os.path.join( - self._temporary_directory.name, - "broken.sock", - ) - os.symlink(missing_target, broken_link) - - _remove_socket_path(broken_link) - - self.assertFalse(os.path.lexists(broken_link)) - - def test_scalar_operand_template_instantiates(self): - operands = [ - _tile_spec(), - {"kind": "scalar", "dtype": "f32", "value": 1.0}, - _tile_spec(), - ] - - mlir = self.client.instantiate("a5", "pto.tadds", operands) - - self.assertIn("func.func @template_tadds", mlir) - self.assertIn("pto.vadds", mlir) - - def test_render_passes_context_attributes_into_template_body(self): - operands = [ - _tile_spec(dtype="f32", shape=(8, 64)), - _tile_spec(dtype="f32", shape=(8, 64)), - _tile_spec(dtype="i8", shape=(8, 64)), - ] - - mlir = self.client.instantiate( - "a5", - "pto.tcmp", - operands, - context_attrs={"cmp_mode": "gt"}, - candidate_id="template_tcmp", - ) - - self.assertIn('"gt"', mlir) - self.assertNotIn('"eq"', mlir) - - def test_vector_operand_metadata_is_accepted(self): - operands = [ - _tile_spec(), - _tile_spec(), - _tile_spec(), - _tile_spec(), - {"kind": "vector", "dtype": "i16", "shape": [4]}, - ] - - metadata = self.client.get_metadata("a5", "pto.tmrgsort", operands) - - self.assertIn("template_tmrgsort_multi_list2", metadata["candidates"]) - - def test_view_operand_template_instantiates(self): - operands = [_view_spec(), _tile_spec()] - - mlir = self.client.instantiate( - "a5", - "pto.tload", - operands, - candidate_id="template_tload_nd2nd", - ) - - self.assertIn("func.func @template_tload_nd2nd", mlir) - self.assertIn("pto.tensor_view_addr", mlir) - self.assertIn("pto.mte_gm_ub", mlir) - - def test_unsupported_operand_kind_is_rejected_explicitly(self): - operands = list(TADD_OPERANDS) - operands[0] = {"kind": "mystery", "dtype": "f32", "shape": [64]} - - with self.assertRaisesRegex(DaemonError, "supports tile, scalar, view, and vector operands"): - self.client.instantiate( - "a5", - "pto.tadd", - operands, - candidate_id=TADD, - ) - - def test_unknown_op_errors(self): - with self.assertRaises(DaemonError): - self.client.instantiate("a5", "pto.tnope", TADD_OPERANDS) - - -if __name__ == "__main__": - unittest.main() diff --git a/ptodsl/tests/test_tilelib_render.py b/ptodsl/tests/test_tilelib_render.py index ca5d4a677f..bc69124e03 100644 --- a/ptodsl/tests/test_tilelib_render.py +++ b/ptodsl/tests/test_tilelib_render.py @@ -12,9 +12,13 @@ (ptodsl differs in SSA naming, constant hoisting, index-vs-i32 carry, ptr typing). """ +import json import unittest from pathlib import Path +from ptoas.mlir.dialects import pto as pto_dialect +from ptoas.mlir.ir import Context +from ptodsl.tilelib._compiler_runtime import materialize from ptodsl.tilelib import TileSpec, f32 from TileOps.a5.tadd import template_tadd @@ -44,6 +48,25 @@ def _render(): return template_tadd.specialize(src0=spec, src1=spec, dst=spec).mlir_text() +def _materialize(context): + tile_spec = { + "kind": "tile", + "shape": [8, 64], + "valid_shape": [8, 64], + "dtype": "f32", + "memory_space": "ub", + "config": {"b_layout": "row_major", "s_layout": "none_box"}, + } + return materialize( + "a5", + "pto.tadd", + json.dumps([tile_spec, tile_spec, tile_spec]), + "{}", + "template_tadd", + context, + ) + + class TileLibRenderTest(unittest.TestCase): def test_renders_structured_abstraction(self): text = _render() @@ -65,6 +88,30 @@ def test_golden_fixture_uses_same_abstraction(self): for op in ("pto.tile_buf_addr", "!pto.ptr", "pto.vlds", "pto.vadd", "pto.vsts", "pto.plt_b32"): self.assertIn(op, golden) + def test_materialize_uses_borrowed_context_and_returns_fresh_modules(self): + context = Context() + pto_dialect.register_dialect(context, load=True) + first, first_entry = _materialize(context) + second, second_entry = _materialize(context) + + self.assertIs(first.context, context) + self.assertIs(second.context, context) + self.assertIsNot(first, second) + self.assertEqual(first_entry, "template_tadd") + self.assertEqual(second_entry, "template_tadd") + self.assertTrue(first.operation.verify()) + self.assertTrue(second.operation.verify()) + self.assertIn("func.func @template_tadd", str(first)) + + def test_materialized_surface_wrappers_release_without_cycle_collection(self): + context = Context() + pto_dialect.register_dialect(context, load=True) + module, _ = _materialize(context) + self.assertEqual(context._get_live_operation_count(), 0) + + del module + self.assertEqual(context._get_live_operation_count(), 0) + if __name__ == "__main__": unittest.main() diff --git a/test/lit/vpto/expand_tile_op_ptodsl_tsub.pto b/test/lit/vpto/expand_tile_op_ptodsl_tsub.pto index 76819a40ee..b7c6b5374a 100644 --- a/test/lit/vpto/expand_tile_op_ptodsl_tsub.pto +++ b/test/lit/vpto/expand_tile_op_ptodsl_tsub.pto @@ -6,10 +6,10 @@ // 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. -// Test that PTOAS can select the PTODSL TileLib daemon and expand the -// single-candidate pto.tsub template without using the legacy TileLang path. +// Test that PTOAS can use the in-process PTODSL TileLib service to expand the +// single-candidate pto.tsub template. // -// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --tile-lib-backend=ptodsl %s -o - 2>/dev/null | FileCheck %s +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto %s -o - 2>/dev/null | FileCheck %s // CHECK: func.func @TSUB // CHECK-NOT: pto.tsub ins diff --git a/test/lit/vpto/expand_tile_op_ptodsl_view_stride_cache.pto b/test/lit/vpto/expand_tile_op_ptodsl_view_stride_cache.pto index 787b00ea8c..963d44f28d 100644 --- a/test/lit/vpto/expand_tile_op_ptodsl_view_stride_cache.pto +++ b/test/lit/vpto/expand_tile_op_ptodsl_view_stride_cache.pto @@ -10,7 +10,7 @@ // bodies. Two tstores with the same tile type but different destination view // strides must therefore not share one cached helper specialization. // -// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --tile-lib-backend=ptodsl %s -o - 2>/dev/null | FileCheck %s +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto %s -o - 2>/dev/null | FileCheck %s // CHECK-LABEL: func.func @STORE_COMPACT // CHECK: %[[COMPACT_GM:.*]] = arith.constant 4 : i64 diff --git a/test/tilelang_st/npu/a5/src/st/smoke/testcase/run_ptoas_to_file.cmake b/test/tilelang_st/npu/a5/src/st/smoke/testcase/run_ptoas_to_file.cmake index 31b9b92438..b8a3a0070a 100644 --- a/test/tilelang_st/npu/a5/src/st/smoke/testcase/run_ptoas_to_file.cmake +++ b/test/tilelang_st/npu/a5/src/st/smoke/testcase/run_ptoas_to_file.cmake @@ -28,17 +28,6 @@ endif() list(APPEND PTOAS_COMMAND --pto-backend=vpto) -set(PTOAS_TILE_LIB_BACKEND_EFFECTIVE "") -if(DEFINED PTOAS_TILE_LIB_BACKEND AND NOT PTOAS_TILE_LIB_BACKEND STREQUAL "") - set(PTOAS_TILE_LIB_BACKEND_EFFECTIVE "${PTOAS_TILE_LIB_BACKEND}") -elseif(DEFINED ENV{PTOAS_TILE_LIB_BACKEND} AND NOT "$ENV{PTOAS_TILE_LIB_BACKEND}" STREQUAL "") - set(PTOAS_TILE_LIB_BACKEND_EFFECTIVE "$ENV{PTOAS_TILE_LIB_BACKEND}") -endif() - -if(NOT PTOAS_TILE_LIB_BACKEND_EFFECTIVE STREQUAL "") - list(APPEND PTOAS_COMMAND "--tile-lib-backend=${PTOAS_TILE_LIB_BACKEND_EFFECTIVE}") -endif() - if(PTOAS_ENABLE_INSERT_SYNC) list(APPEND PTOAS_COMMAND --enable-insert-sync) endif() diff --git a/test/tilelang_st/npu/a5/src/st/testcase/run_ptoas_to_file.cmake b/test/tilelang_st/npu/a5/src/st/testcase/run_ptoas_to_file.cmake index 31b9b92438..b8a3a0070a 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/run_ptoas_to_file.cmake +++ b/test/tilelang_st/npu/a5/src/st/testcase/run_ptoas_to_file.cmake @@ -28,17 +28,6 @@ endif() list(APPEND PTOAS_COMMAND --pto-backend=vpto) -set(PTOAS_TILE_LIB_BACKEND_EFFECTIVE "") -if(DEFINED PTOAS_TILE_LIB_BACKEND AND NOT PTOAS_TILE_LIB_BACKEND STREQUAL "") - set(PTOAS_TILE_LIB_BACKEND_EFFECTIVE "${PTOAS_TILE_LIB_BACKEND}") -elseif(DEFINED ENV{PTOAS_TILE_LIB_BACKEND} AND NOT "$ENV{PTOAS_TILE_LIB_BACKEND}" STREQUAL "") - set(PTOAS_TILE_LIB_BACKEND_EFFECTIVE "$ENV{PTOAS_TILE_LIB_BACKEND}") -endif() - -if(NOT PTOAS_TILE_LIB_BACKEND_EFFECTIVE STREQUAL "") - list(APPEND PTOAS_COMMAND "--tile-lib-backend=${PTOAS_TILE_LIB_BACKEND_EFFECTIVE}") -endif() - if(PTOAS_ENABLE_INSERT_SYNC) list(APPEND PTOAS_COMMAND --enable-insert-sync) endif() diff --git a/test/tilelang_st/script/run_a5_st_all_parallel.py b/test/tilelang_st/script/run_a5_st_all_parallel.py index 8299c743f5..d4fc1cd4fb 100755 --- a/test/tilelang_st/script/run_a5_st_all_parallel.py +++ b/test/tilelang_st/script/run_a5_st_all_parallel.py @@ -118,15 +118,11 @@ def _run_one(job, args, ptoas_bin, output_root, base_env): build_dir = job_root / "build" tmp_dir = job_root / "tmp" log_path = output_root / "logs" / f"{job_name}.log" - socket_path = Path("/tmp") / f"ptoas_st_{kind}_{testcase}_{os.getpid()}.sock" started = time.time() env = base_env.copy() env["TMPDIR"] = str(tmp_dir) env["PTODSL_CACHE_DIR"] = str(job_root / "ptodsl-cache") - env["PTOAS_DAEMON_SOCKET_PATH"] = str(socket_path) - if args.tile_lib_backend: - env["PTOAS_TILE_LIB_BACKEND"] = args.tile_lib_backend tmp_dir.mkdir(parents=True, exist_ok=True) (output_root / "logs").mkdir(parents=True, exist_ok=True) @@ -139,69 +135,56 @@ def _run_one(job, args, ptoas_bin, output_root, base_env): "seconds": 0.0, "log": str(log_path), "build_dir": str(build_dir), - "socket": str(socket_path), } - try: - with log_path.open("w", encoding="utf-8") as log_handle: - log_handle.write(f"# kind: {kind}\n") - log_handle.write(f"# testcase: {testcase}\n") - log_handle.write(f"# source: {job['target_dir']}\n") - log_handle.write(f"# build: {build_dir}\n") - log_handle.write(f"# PTOAS_DAEMON_SOCKET_PATH={socket_path}\n") - log_handle.write(f"# PTODSL_CACHE_DIR={env['PTODSL_CACHE_DIR']}\n") - if args.tile_lib_backend: - log_handle.write(f"# PTOAS_TILE_LIB_BACKEND={args.tile_lib_backend}\n") - log_handle.write("\n") - - cmake_cmd = [ - "cmake", - "-S", - job["target_dir"], - "-B", - build_dir, - f"-DRUN_MODE={args.run_mode}", - f"-DSOC_VERSION={DEFAULT_SOC_VERSION}", - f"-DTEST_CASE={testcase}", - f"-DPTOAS_BIN={ptoas_bin}", - f"-DPTOAS_DAEMON_SOCKET_PATH={socket_path}", - ] - if args.tile_lib_backend: - cmake_cmd.append(f"-DPTOAS_TILE_LIB_BACKEND={args.tile_lib_backend}") - - rc = _run_logged(cmake_cmd, log_handle, output_root, env) - if rc == 0: - rc = _run_logged( - ["cmake", "--build", build_dir, "--parallel", str(args.build_jobs)], - log_handle, - output_root, - env, - ) + with log_path.open("w", encoding="utf-8") as log_handle: + log_handle.write(f"# kind: {kind}\n") + log_handle.write(f"# testcase: {testcase}\n") + log_handle.write(f"# source: {job['target_dir']}\n") + log_handle.write(f"# build: {build_dir}\n") + log_handle.write(f"# PTODSL_CACHE_DIR={env['PTODSL_CACHE_DIR']}\n") + log_handle.write("\n") + + cmake_cmd = [ + "cmake", + "-S", + job["target_dir"], + "-B", + build_dir, + f"-DRUN_MODE={args.run_mode}", + f"-DSOC_VERSION={DEFAULT_SOC_VERSION}", + f"-DTEST_CASE={testcase}", + f"-DPTOAS_BIN={ptoas_bin}", + ] + + rc = _run_logged(cmake_cmd, log_handle, output_root, env) + if rc == 0: + rc = _run_logged( + ["cmake", "--build", build_dir, "--parallel", str(args.build_jobs)], + log_handle, + output_root, + env, + ) + if rc != 0: + result["returncode"] = rc + result["phase"] = "build" + result["seconds"] = time.time() - started + return result + + case_work_dir = build_dir / "testcase" / testcase + _copy_case_scripts(job["testcase_root"], testcase, case_work_dir) + + for phase, command in ( + ("gen_data", [sys.executable, "gen_data.py"]), + ("run", [build_dir / "bin" / testcase]), + ("compare", [sys.executable, "compare.py"]), + ): + rc = _run_logged(command, log_handle, case_work_dir, env) if rc != 0: result["returncode"] = rc - result["phase"] = "build" - result["seconds"] = time.time() - started - return result - - case_work_dir = build_dir / "testcase" / testcase - _copy_case_scripts(job["testcase_root"], testcase, case_work_dir) - - for phase, command in ( - ("gen_data", [sys.executable, "gen_data.py"]), - ("run", [build_dir / "bin" / testcase]), - ("compare", [sys.executable, "compare.py"]), - ): - rc = _run_logged(command, log_handle, case_work_dir, env) - if rc != 0: - result["returncode"] = rc - result["phase"] = phase - break - finally: - try: - socket_path.unlink(missing_ok=True) - except OSError: - pass + result["phase"] = phase + break result["seconds"] = time.time() - started return result @@ -270,11 +253,6 @@ def _parse_args(): default=str(_default_output_root(repo_root)), help="Directory for logs, summaries, and per-testcase build trees.", ) - parser.add_argument( - "--tile-lib-backend", - default=os.environ.get("PTOAS_TILE_LIB_BACKEND", ""), - help="Optional PTOAS tile-lib backend, for example ptodsl.", - ) parser.add_argument("--full-only", action="store_true", help="Run only non-smoke ST cases.") parser.add_argument("--smoke-only", action="store_true", help="Run only smoke ST cases.") parser.add_argument("--list", action="store_true", help="List selected jobs and exit.") @@ -339,9 +317,7 @@ def main(): print(f"[INFO] run_mode={args.run_mode} soc={SOC_VERSION} ({DEFAULT_SOC_VERSION})") print(f"[INFO] ptoas={ptoas_bin}") print(f"[INFO] output_root={output_root}") - if args.tile_lib_backend: - print(f"[INFO] PTOAS_TILE_LIB_BACKEND={args.tile_lib_backend}") - print("[INFO] each testcase uses its own build dir, PTODSL cache, TMPDIR, and daemon socket") + print("[INFO] each testcase uses its own build dir, PTODSL cache, and TMPDIR") results = [] max_workers = min(args.jobs, len(jobs)) diff --git a/test/tilelang_st/script/run_ptodsl_st_parallel.py b/test/tilelang_st/script/run_ptodsl_st_parallel.py index 4c1e4a0ac7..045cb94ab8 100755 --- a/test/tilelang_st/script/run_ptodsl_st_parallel.py +++ b/test/tilelang_st/script/run_ptodsl_st_parallel.py @@ -7,7 +7,7 @@ # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. -"""Run TileLang ST testcases with the PTODSL TileLib backend and per-test logs.""" +"""Run TileLang ST testcases with PTODSL TileLib and per-test logs.""" import argparse import concurrent.futures @@ -69,7 +69,6 @@ def _run_build(args, default_soc_version, target_dir, log_dir, ptoas_bin): started = time.time() with build_log.open("w", encoding="utf-8") as handle: handle.write(f"# cwd: {target_dir}\n") - handle.write(f"# PTOAS_TILE_LIB_BACKEND={args.tile_lib_backend}\n") handle.write(f"# ptoas: {ptoas_bin}\n") handle.write("# command: run_st.build_project(..., testcase='all', ...)\n\n") handle.flush() @@ -116,13 +115,11 @@ def _run_one_testcase(args, testcase, target_dir, log_dir, ptoas_bin): command.append("--smoke") env = os.environ.copy() - env["PTOAS_TILE_LIB_BACKEND"] = args.tile_lib_backend started = time.time() with log_path.open("w", encoding="utf-8") as handle: handle.write(f"# testcase: {testcase}\n") handle.write(f"# cwd: {target_dir}\n") - handle.write(f"# PTOAS_TILE_LIB_BACKEND={args.tile_lib_backend}\n") handle.write("# command: " + " ".join(command) + "\n\n") handle.flush() proc = subprocess.Popen( @@ -150,10 +147,7 @@ def _run_one_testcase(args, testcase, target_dir, log_dir, ptoas_bin): def _parse_args(): repo_root = _repo_root() parser = argparse.ArgumentParser( - description=( - "Run TileLang ST testcases in parallel with PTOAS_TILE_LIB_BACKEND=ptodsl " - "and save one log per testcase." - ) + description="Run TileLang ST testcases in parallel and save one log per testcase." ) parser.add_argument("-r", "--run-mode", default="sim", help="Run mode: sim or npu.") parser.add_argument("-v", "--soc-version", default="a5", help="SoC version key, default: a5.") @@ -182,11 +176,6 @@ def _parse_args(): default=None, help="Directory for build.log, one .log per testcase, and summary files.", ) - parser.add_argument( - "--tile-lib-backend", - default="ptodsl", - help="Value for PTOAS_TILE_LIB_BACKEND, default: ptodsl.", - ) parser.add_argument( "--full", action="store_true", @@ -250,14 +239,12 @@ def main(): print(f"[INFO] target_dir={target_dir}") print(f"[INFO] ptoas={ptoas_bin}") print(f"[INFO] logs={log_dir}") - print(f"[INFO] PTOAS_TILE_LIB_BACKEND={args.tile_lib_backend}") if args.dry_run: for testcase in selected: print(f"[DRY-RUN] {testcase}") return 0 - os.environ["PTOAS_TILE_LIB_BACKEND"] = args.tile_lib_backend default_soc_version = run_all_st.SOC_VERSION_MAP[args.soc_version] results = [] @@ -311,7 +298,6 @@ def main(): break summary = { - "backend": args.tile_lib_backend, "run_mode": args.run_mode, "soc_version": args.soc_version, "smoke": args.smoke, diff --git a/tools/ptoas/CMakeLists.txt b/tools/ptoas/CMakeLists.txt index 2ce4929432..d2aadd3c00 100644 --- a/tools/ptoas/CMakeLists.txt +++ b/tools/ptoas/CMakeLists.txt @@ -16,7 +16,6 @@ set(PTOAS_RUNTIME_SOURCES driver.cpp VPTOHostStubEmission.cpp ObjectEmission.cpp - TilelangDaemon.cpp ) add_library(PTOASVFSIMTSizePatcher STATIC @@ -67,13 +66,6 @@ endforeach() function(ptoas_configure_runtime_compile_target target_name) target_compile_definitions(${target_name} PRIVATE PTOAS_RELEASE_VERSION="${PTOAS_CLI_VERSION}" - # Source-tree defaults for TileLib expansion. These let ptoas run directly - # from the build tree without passing --ptodsl-pkg-path / - # --tileops-pkg-path. Installed layouts are expected - # to launch through the wrapper/launcher flow, which injects explicit - # runtime paths instead of relying on executable-path probing. - PTOAS_DEFAULT_PTODSL_PKG_PATH="${CMAKE_SOURCE_DIR}/ptodsl" - PTOAS_DEFAULT_TILEOPS_PKG_PATH="${CMAKE_SOURCE_DIR}/lib" ${ARGN} ) add_dependencies(${target_name} diff --git a/tools/ptoas/NativeModule.cpp b/tools/ptoas/NativeModule.cpp index 5a474cbd27..0f0bda26d4 100644 --- a/tools/ptoas/NativeModule.cpp +++ b/tools/ptoas/NativeModule.cpp @@ -9,7 +9,12 @@ #include "ptoas.h" #include "PTOModule.h" +#include "PTO/Transforms/TileLibService.h" +#include "mlir/Bindings/Python/PybindAdaptors.h" +#include "mlir/CAPI/IR.h" + +#include "llvm/Support/raw_ostream.h" #include "pybind11/pybind11.h" #include "pybind11/stl.h" @@ -20,6 +25,112 @@ namespace py = pybind11; namespace { +class PythonTileLibService final : public mlir::pto::TileLibService { +public: + mlir::FailureOr + getMetadata(const mlir::pto::TileLibMaterializationRequest &request) override { + py::gil_scoped_acquire acquire; + try { + return py::cast(getCompilerRuntime().attr("metadata")( + request.target, request.op, request.operandSpecsJson, + request.contextAttrsJson)); + } catch (const py::error_already_set &error) { + llvm::errs() << "TileLib: PTODSL metadata query raised Python " + "exception:\n" + << error.what() << "\n"; + return mlir::failure(); + } + } + + mlir::LogicalResult + materialize(const mlir::pto::TileLibMaterializationRequest &request, + mlir::MLIRContext &context, + mlir::pto::TileLibMaterializationCallback callback) override { + py::gil_scoped_acquire acquire; + try { + py::object contextOwner = getPythonContext(context); + MlirContext pythonContext = py::cast(contextOwner); + if (unwrap(pythonContext) != &context) { + llvm::errs() << "TileLib: Python context does not match the PTOAS " + "MLIRContext\n"; + return mlir::failure(); + } + + py::tuple result = getCompilerRuntime().attr("materialize")( + request.target, request.op, request.operandSpecsJson, + request.contextAttrsJson, request.candidateId, contextOwner); + if (result.size() != 2) + throw py::value_error( + "PTODSL materialize() must return (module, entry_symbol)"); + + // MlirModule is a non-owning handle. Keep result[0] alive until the + // complete source module has been cloned into C++ ownership. + py::object moduleOwner = result[0]; + MlirModule rawModule = py::cast(moduleOwner); + if (!mlirContextEqual(mlirModuleGetContext(rawModule), pythonContext)) { + llvm::errs() << "TileLib: PTODSL returned a module from a different " + "MLIRContext\n"; + return mlir::failure(); + } + + mlir::ModuleOp source = unwrap(rawModule); + return callback(source, py::cast(result[1])); + } catch (const py::error_already_set &error) { + llvm::errs() << "TileLib: PTODSL materialization raised Python " + "exception:\n" + << error.what() << "\n"; + return mlir::failure(); + } catch (const std::exception &error) { + llvm::errs() << "TileLib: invalid PTODSL materialization result: " + << error.what() << "\n"; + return mlir::failure(); + } + } + +private: + static py::module_ getCompilerRuntime() { + // Python's sys.modules cache makes this a process-wide runtime module + // without storing a py::object whose destructor could outlive CPython. + return py::module_::import("ptodsl.tilelib._compiler_runtime"); + } + + static py::object getPythonContext(mlir::MLIRContext &context) { + py::object capsule = py::reinterpret_steal( + mlirPythonContextToCapsule(wrap(&context))); + return py::module_::import("ptoas.mlir.ir") + .attr("Context") + .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule); + } +}; + +constexpr char kRuntimeRegistrationCapsuleName[] = + "ptoas.TileLibRuntimeRegistration"; + +class PythonTileLibRuntimeRegistration { +public: + PythonTileLibRuntimeRegistration() + : service(std::make_shared()) { + mlir::pto::TileLibRuntime::install(service); + } + + ~PythonTileLibRuntimeRegistration() { + mlir::pto::TileLibRuntime::uninstall(service.get()); + } + +private: + std::shared_ptr service; +}; + +void destroyRuntimeRegistration(PyObject *capsule) { + void *pointer = + PyCapsule_GetPointer(capsule, kRuntimeRegistrationCapsuleName); + if (!pointer) { + PyErr_Clear(); + return; + } + delete static_cast(pointer); +} + int runPTOASFromPython(const std::vector &arguments) { std::vector storage = arguments; std::vector argv; @@ -27,8 +138,17 @@ int runPTOASFromPython(const std::vector &arguments) { for (std::string &argument : storage) argv.push_back(argument.data()); - py::gil_scoped_release release; - return mlir::pto::runPTOAS(static_cast(argv.size()), argv.data()); + py::object contextOwner = + py::module_::import("ptoas.mlir.ir").attr("Context")(); + MlirContext rawContext = py::cast(contextOwner); + + int result; + { + py::gil_scoped_release release; + result = mlir::pto::runPTOAS(static_cast(argv.size()), argv.data(), + *unwrap(rawContext)); + } + return result; } } // namespace @@ -37,5 +157,10 @@ PYBIND11_MODULE(_core, module) { module.doc() = "PTOAS compiler and PTO dialect native bindings"; py::module_::import("ptoas.mlir.ir"); mlir::pto::python::populatePTODialectBindings(module); + module.add_object( + "_tilelib_runtime_registration", + py::capsule(new PythonTileLibRuntimeRegistration(), + kRuntimeRegistrationCapsuleName, + destroyRuntimeRegistration)); module.def("main", &runPTOASFromPython, py::arg("argv")); } diff --git a/tools/ptoas/TilelangDaemon.cpp b/tools/ptoas/TilelangDaemon.cpp deleted file mode 100644 index 9ccc86221c..0000000000 --- a/tools/ptoas/TilelangDaemon.cpp +++ /dev/null @@ -1,154 +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 "PTO/Support/PythonExecutable.h" -#include "TilelangDaemon.h" -#include "llvm/ADT/SmallVector.h" -#include "llvm/ADT/StringRef.h" -#include "llvm/Support/FileSystem.h" -#include "llvm/Support/Program.h" -#include -#include -#include -#include -#include -#include - -extern char **environ; - -namespace ptoas { - -std::optional> DaemonManager::processInfo; - -std::string DaemonManager::generateSocketPath() { - return "/tmp/tilelib_daemon_" + std::to_string(::getpid()) + ".sock"; -} - -bool DaemonManager::start(const std::string &socketPath, - const std::string &daemonModule, - const std::string &pythonExe, - const std::string &pkgPath, - const std::string &templateDir) { - auto pythonPath = - mlir::pto::resolvePythonExecutable(pythonExe.empty() ? "python3" - : pythonExe); - if (!pythonPath) { - llvm::errs() << "Error: Cannot find Python executable '" - << (pythonExe.empty() ? "python3" : pythonExe) - << "' for daemon\n"; - return false; - } - - llvm::SmallVector args = { - *pythonPath, "-m", daemonModule, "--socket", socketPath, - }; - if (!templateDir.empty()) { - args.push_back("--template-dir"); - args.push_back(templateDir); - } - - llvm::SmallVector envp; - std::string pythonPathEnv; - std::vector envStorage; - - if (!pkgPath.empty()) { - const char *existingPath = ::getenv("PYTHONPATH"); - pythonPathEnv = "PYTHONPATH=" + pkgPath; - if (existingPath && existingPath[0] != '\0') { - pythonPathEnv += ":"; - pythonPathEnv += existingPath; - } - for (char **e = environ; *e; ++e) { - llvm::StringRef entry(*e); - if (entry.starts_with("PYTHONPATH=")) - continue; - envStorage.push_back(std::string(entry)); - } - envStorage.push_back(pythonPathEnv); - for (auto &s : envStorage) - envp.push_back(s); - } - - std::string errMsg; - bool executionFailed = false; - - llvm::sys::ProcessInfo procInfo = llvm::sys::ExecuteNoWait( - *pythonPath, args, - !pkgPath.empty() - ? std::optional>(envp) - : std::nullopt, - {}, 0, &errMsg, &executionFailed, nullptr, true); - - if (executionFailed || procInfo.Pid == llvm::sys::ProcessInfo::InvalidPid) { - llvm::errs() << "Error: Failed to start TileLib daemon module '" - << daemonModule << "': " << errMsg << "\n"; - return false; - } - - processInfo = std::make_pair(procInfo.Pid, socketPath); - - // Python startup time depends on the selected TileLib frontend and its - // imports. Poll instead of relying on one fixed sleep. - bool socketReady = false; - // PTODSL imports can be noticeably slower on heavily loaded CI runners where - // many ptoas processes start TileLib daemons concurrently. - for (int attempt = 0; attempt < 600; ++attempt) { - std::this_thread::sleep_for(std::chrono::milliseconds(50)); - if (llvm::sys::fs::exists(socketPath)) { - socketReady = true; - break; - } - } - - if (!socketReady) { - llvm::errs() << "Error: Daemon socket not created at " << socketPath << "\n"; - llvm::errs() << "Note: Daemon process started (pid=" << procInfo.Pid - << ") but socket not found. Check daemon logs.\n"; - kill(procInfo.Pid, SIGTERM); - processInfo = std::nullopt; - return false; - } - - llvm::errs() << "TileLib daemon '" << daemonModule << "' started (pid=" - << procInfo.Pid - << ", socket=" << socketPath << ")\n"; - return true; -} - -void DaemonManager::stop() { - if (!processInfo) - return; - - int pid = processInfo->first; - std::string socketPath = processInfo->second; - - kill(pid, SIGTERM); - - std::this_thread::sleep_for(std::chrono::milliseconds(100)); - - if (llvm::sys::fs::exists(socketPath)) { - llvm::sys::fs::remove(socketPath); - } - - llvm::errs() << "TileLib daemon stopped (pid=" << pid << ")\n"; - processInfo = std::nullopt; -} - -bool DaemonManager::isRunning() { - return processInfo.has_value(); -} - -static void daemonCleanupHandler() { - DaemonManager::stop(); -} - -void registerDaemonCleanup() { - std::atexit(daemonCleanupHandler); -} - -} // namespace ptoas diff --git a/tools/ptoas/TilelangDaemon.h b/tools/ptoas/TilelangDaemon.h deleted file mode 100644 index 8b369f6fba..0000000000 --- a/tools/ptoas/TilelangDaemon.h +++ /dev/null @@ -1,44 +0,0 @@ -// Copyright (c) 2026 Huawei Technologies Co., Ltd. -// This program is free software, you can redistribute it and/or modify it under the terms and conditions of -// CANN Open Software License Agreement Version 2.0 (the "License"). -// Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -// See LICENSE in the root of the software repository for the full text of the License. - -#ifndef PTOAS_TILELANG_DAEMON_H -#define PTOAS_TILELANG_DAEMON_H - -#include -#include -#include - -namespace llvm::sys { -using procid_t = int; -} - -namespace ptoas { - -class DaemonManager { -public: - static std::string generateSocketPath(); - - static bool start(const std::string &socketPath, - const std::string &daemonModule, - const std::string &pythonExe, - const std::string &pkgPath, - const std::string &templateDir = ""); - - static void stop(); - - static bool isRunning(); - -private: - static std::optional> processInfo; -}; - -void registerDaemonCleanup(); - -} // namespace ptoas - -#endif // PTOAS_TILELANG_DAEMON_H diff --git a/tools/ptoas/driver.cpp b/tools/ptoas/driver.cpp index 4c44d980ce..d9b06c4249 100644 --- a/tools/ptoas/driver.cpp +++ b/tools/ptoas/driver.cpp @@ -678,7 +678,14 @@ static LogicalResult emitVPTOLLVMFatobj( mlir::pto::PTOASContext::PTOASContext(DialectRegistry ®istry, llvm::StringRef outputPath, int argc, char **argv) - : mlirContext(registry), outputPath(outputPath.str()), argc(argc), + : ownedMlirContext(std::make_unique(registry)), + mlirContext(ownedMlirContext.get()), outputPath(outputPath.str()), + argc(argc), argv(argv) {} + +mlir::pto::PTOASContext::PTOASContext( + MLIRContext &borrowedContext, llvm::StringRef outputPath, int argc, + char **argv) + : mlirContext(&borrowedContext), outputPath(outputPath.str()), argc(argc), argv(argv) {} mlir::pto::PTOASContext::~PTOASContext() = default; @@ -694,11 +701,11 @@ mlir::pto::PTOASContext::initializeEnvironment(bool requiresToolchain, void mlir::pto::PTOASContext::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); + mlirContext->allowUnregisteredDialects(true); + mlir::pto::loadPTOASDialects(*mlirContext); } -MLIRContext &mlir::pto::PTOASContext::getMLIRContext() { return mlirContext; } +MLIRContext &mlir::pto::PTOASContext::getMLIRContext() { return *mlirContext; } void mlir::pto::PTOASContext::setArch(std::string value) { arch = std::move(value); @@ -1266,9 +1273,12 @@ static LogicalResult writeTextOutput(llvm::StringRef output, // +-------------+ +------------------------------------------+ // | C++ source | | fatobj | // +-------------+ +------------------------------------------+ -static int runPTOASDriver(int argc, char **argv) { +static int runPTOASDriver(int argc, char **argv, + MLIRContext *borrowedContext = nullptr) { DialectRegistry registry; mlir::pto::registerPTOASDialects(registry); + if (borrowedContext) + borrowedContext->appendDialectRegistry(registry); mlir::pto::registerPTOASPassesAndCLOptions(); llvm::cl::SetVersionPrinter(printPTOASVersion); @@ -1283,10 +1293,17 @@ static int runPTOASDriver(int argc, char **argv) { llvm::errs())) return 1; - PTOASContext context(registry, outputFilename, argc, argv); - context.setOutputCANNVersionOverride(outputCANNVersionOverride); - context.setVFSIMTSizeFixMode(mlir::pto::vptoFixVFSIMTSize); - context.initializeMLIRContext(); + std::unique_ptr context; + if (borrowedContext) { + context = std::make_unique(*borrowedContext, outputFilename, + argc, argv); + } else { + context = + std::make_unique(registry, outputFilename, argc, argv); + } + context->setOutputCANNVersionOverride(outputCANNVersionOverride); + context->setVFSIMTSizeFixMode(mlir::pto::vptoFixVFSIMTSize); + context->initializeMLIRContext(); std::unique_ptr inputBuffer = readInputBuffer(); if (!inputBuffer) @@ -1294,24 +1311,24 @@ static int runPTOASDriver(int argc, char **argv) { std::string arch; OwningOpRef module = loadInputModule( - std::move(inputBuffer), context.getMLIRContext(), cliArchSpecified, arch); + std::move(inputBuffer), context->getMLIRContext(), cliArchSpecified, arch); if (!module) return 1; - context.setArch(std::move(arch)); + context->setArch(std::move(arch)); mlir::pto::BackendInfo backendInfo; if (failed(buildBackendInfo(module.get(), cliBackendSpecified, backendInfo))) return 1; - context.setBackendInfo(std::move(backendInfo)); - (void)context.initializeEnvironment(context.getBackendInfo().requiresToolchain, - llvm::errs()); + context->setBackendInfo(std::move(backendInfo)); + (void)context->initializeEnvironment( + context->getBackendInfo().requiresToolchain, llvm::errs()); mlir::pto::PTOASCompileResult result; - if (failed(runPTOASJobs(module, context, result))) + if (failed(runPTOASJobs(module, *context, result))) return 1; if (result.kind == mlir::pto::PTOASCompileResultKind::Text) - return failed(writeTextOutput(result.textOutput, context.getOutputPath())); + return failed(writeTextOutput(result.textOutput, context->getOutputPath())); if (result.kind == mlir::pto::PTOASCompileResultKind::MixedObject) return 0; @@ -1322,3 +1339,8 @@ static int runPTOASDriver(int argc, char **argv) { int mlir::pto::runPTOAS(int argc, char **argv) { return runPTOASDriver(argc, argv); } + +int mlir::pto::runPTOAS(int argc, char **argv, + MLIRContext &borrowedContext) { + return runPTOASDriver(argc, argv, &borrowedContext); +} diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index 654ef4ed96..6c795c3d10 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -14,7 +14,6 @@ #include "PTO/Transforms/Passes.h" #include "PTO/Transforms/BufferizableOpInterfaceImpl.h" #include "VPTOHostStubEmission.h" -#include "TilelangDaemon.h" #include "PTO/Transforms/CppPostprocess.h" #include "mlir/AsmParser/AsmParserState.h" #include "mlir/IR/MLIRContext.h" @@ -214,7 +213,6 @@ void mlir::pto::registerPTOASPassesAndCLOptions() { mlir::pto::registerPTOPasses(); mlir::pto::registerPTOInlineLibCall(); mlir::pto::registerFoldTileBufIntrinsics(); - mlir::pto::registerExpandTileOp(); mlir::pto::registerLowerPTOToUBufOps(); mlir::registerPassManagerCLOptions(); } @@ -230,48 +228,6 @@ void mlir::pto::loadPTOASDialects(MLIRContext &context) { context.getOrLoadDialect(); } -static bool pathExists(llvm::StringRef path) { - return !path.empty() && llvm::sys::fs::exists(path); -} - -static std::string getParentDir(llvm::StringRef path) { - llvm::SmallString<256> parent(path); - llvm::sys::path::remove_filename(parent); - llvm::sys::path::remove_dots(parent, true); - return std::string(parent); -} - -static std::string joinPath(llvm::StringRef lhs, llvm::StringRef rhs) { - llvm::SmallString<256> joined(lhs); - llvm::sys::path::append(joined, rhs); - llvm::sys::path::remove_dots(joined, true); - return std::string(joined); -} - -static std::string detectInstalledPythonPkgRoot(const char *argv0, - llvm::StringRef packageName) { - std::string exePath = llvm::sys::fs::getMainExecutable(argv0, (void *)&main); - if (exePath.empty()) - return {}; - - const std::string exeDir = getParentDir(exePath); - const std::string prefixDir = getParentDir(exeDir); - const std::string installedPkg = joinPath(prefixDir, packageName); - if (pathExists(installedPkg)) - return prefixDir; - return {}; -} - -static bool hasCLIOption(int argc, char **argv, llvm::StringRef option) { - const std::string optionWithValue = (option + "=").str(); - for (int i = 1; i < argc; ++i) { - llvm::StringRef arg(argv[i]); - if (arg == option || arg.starts_with(optionWithValue)) - return true; - } - return false; -} - static LogicalResult applyConfiguredPassManagerCLOptions( PassManager &pm, llvm::StringRef pipelineName, llvm::raw_ostream &diagOS = llvm::errs()) { @@ -453,128 +409,6 @@ static llvm::cl::opt enableTileOpExpand( "--pto-backend=vpto."), llvm::cl::init(false)); -#ifndef PTOAS_DEFAULT_PTODSL_PKG_PATH -#define PTOAS_DEFAULT_PTODSL_PKG_PATH "" -#endif -#ifndef PTOAS_DEFAULT_TILEOPS_PKG_PATH -#define PTOAS_DEFAULT_TILEOPS_PKG_PATH "" -#endif - -static llvm::cl::opt ptodslPkgPath( - "ptodsl-pkg-path", - llvm::cl::desc("PYTHONPATH for the ptodsl package " - "(default: /ptodsl, baked in at build time)"), - llvm::cl::init(PTOAS_DEFAULT_PTODSL_PKG_PATH)); - -static llvm::cl::opt tileopsPkgPath( - "tileops-pkg-path", - llvm::cl::desc("PYTHONPATH for the TileOps PTODSL template package " - "(default: /lib, baked in at build time)"), - llvm::cl::init(PTOAS_DEFAULT_TILEOPS_PKG_PATH)); - -static llvm::cl::opt daemonSocketPath( - "daemon-socket-path", - llvm::cl::desc("Path to Unix domain socket for daemon RPC " - "(default: /tmp/tilelib_daemon_{pid}.sock)"), - llvm::cl::init("")); - -enum class TileLibBackend { - PTODSL, -}; - -static llvm::cl::opt tileLibBackend( - "tile-lib-backend", - llvm::cl::desc("TileLib backend used by ExpandTileOp"), - llvm::cl::values( - clEnumValN(TileLibBackend::PTODSL, "ptodsl", - "Use the PTODSL TileLib daemon")), - llvm::cl::init(TileLibBackend::PTODSL)); - -static std::string resolveTileLibPythonExe() { - const char *pythonExe = ::getenv("PTOAS_PYTHON_EXE"); - if (pythonExe && pythonExe[0] != '\0') - return pythonExe; - return "python3"; -} - -static pto::ExpandTileOpOptions resolveExpandTileOpOptions(int argc, - char **argv) { - pto::ExpandTileOpOptions expandOpts; - expandOpts.pythonExe = resolveTileLibPythonExe(); - std::string resolvedPtodslPkgPath = ptodslPkgPath; - std::string resolvedTileOpsPkgPath = tileopsPkgPath; - - if (!hasCLIOption(argc, argv, "--ptodsl-pkg-path")) { - const char *envPtodslRoot = ::getenv("PTODSL_PYTHON_ROOT"); - if (envPtodslRoot && envPtodslRoot[0] != '\0') - resolvedPtodslPkgPath = envPtodslRoot; - else { - std::string installedPtodslPkgPath = - detectInstalledPythonPkgRoot(argv[0], "ptodsl"); - if (!installedPtodslPkgPath.empty()) - resolvedPtodslPkgPath = installedPtodslPkgPath; - } - } - - if (!hasCLIOption(argc, argv, "--tileops-pkg-path")) { - const char *envTileOpsRoot = ::getenv("PTO_TILEOPS_PYTHON_ROOT"); - if (envTileOpsRoot && envTileOpsRoot[0] != '\0') - resolvedTileOpsPkgPath = envTileOpsRoot; - else { - std::string installedTileOpsPkgPath = - detectInstalledPythonPkgRoot(argv[0], "TileOps"); - if (!installedTileOpsPkgPath.empty()) - resolvedTileOpsPkgPath = installedTileOpsPkgPath; - } - } - - expandOpts.tileLibBackend = "ptodsl"; - expandOpts.daemonHelperModule = "ptodsl.tilelib.serving.helper"; - expandOpts.tileLibPkgPath = resolvedPtodslPkgPath; - if (!resolvedTileOpsPkgPath.empty()) { - if (!expandOpts.tileLibPkgPath.empty()) - expandOpts.tileLibPkgPath += ":"; - expandOpts.tileLibPkgPath += resolvedTileOpsPkgPath; - } - - // Daemon mode is default (no CLI option needed) - // Automatically start daemon for instance caching - std::string socket = daemonSocketPath; - if (socket.empty()) - socket = ptoas::DaemonManager::generateSocketPath(); - - // Register cleanup handler (daemon will be stopped on PTOAS exit) - ptoas::registerDaemonCleanup(); - - // Try to start daemon automatically - if (ptoas::DaemonManager::start(socket, "ptodsl.tilelib.serving.daemon", - expandOpts.pythonExe, - expandOpts.tileLibPkgPath, "")) { - expandOpts.daemonSocketPath = socket; - llvm::errs() << "Info: " << expandOpts.tileLibBackend - << " TileLib daemon started successfully\n"; - } else { - expandOpts.daemonSocketPath = ""; - llvm::errs() - << "Error: Failed to start the PTODSL TileLib daemon; no TileLang " - "fallback will be used\n"; - } - - return expandOpts; -} - - -static pto::InsertTemplateAttributesOptions -buildInsertTemplateAttributesOptions( - const pto::ExpandTileOpOptions &expandOptions) { - pto::InsertTemplateAttributesOptions options; - options.pythonExe = expandOptions.pythonExe; - options.daemonSocketPath = expandOptions.daemonSocketPath; - options.tileLibPkgPath = expandOptions.tileLibPkgPath; - options.daemonHelperModule = expandOptions.daemonHelperModule; - return options; -} - static llvm::cl::opt enableOpFusion( "enable-op-fusion", llvm::cl::desc("Control A5 tile fusion on level2/level3. Defaults to " @@ -2951,9 +2785,7 @@ static void prepareVPTOForEmission(PassManager &pm) { kernelModulePM.addPass(pto::createPTOValidateVPTOEmissionIRPass()); } -static void -lowerPTOToVPTOBackend(PassManager &pm, ModuleOp module, - const pto::ExpandTileOpOptions &expandOpts) { +static void lowerPTOToVPTOBackend(PassManager &pm, ModuleOp module) { auto &kernelModulePM = pm.nest(); auto moduleArchAttr = module->getAttrOfType("pto.target_arch"); @@ -2970,7 +2802,7 @@ lowerPTOToVPTOBackend(PassManager &pm, ModuleOp module, return; } - kernelModulePM.addPass(pto::createExpandTileOpPass(expandOpts)); + kernelModulePM.addPass(pto::createExpandTileOpPass()); kernelModulePM.addPass(pto::createPTOInlineLibCallPass()); kernelModulePM.addNestedPass( @@ -3065,21 +2897,13 @@ static int emitVPTOBackendResult(ModuleOp module, PTOASCompileResult &result, } static LogicalResult runVPTOBackendPipeline(OwningOpRef &module, - bool hasTileOpsToExpand, - const pto::ExpandTileOpOptions - *expandOptions) { + bool hasTileOpsToExpand) { PassManager pm(module->getContext()); pm.enableVerifier(); pm.addPass(pto::createVPTOSplitCVModulePass()); pm.addPass(pto::createVPTONormalizeContainerPass()); - if (hasTileOpsToExpand) { - if (!expandOptions) { - llvm::errs() << "Error: tile expansion requires resolved TileLib " - "options.\n"; - return failure(); - } - lowerPTOToVPTOBackend(pm, module.get(), *expandOptions); - } + if (hasTileOpsToExpand) + lowerPTOToVPTOBackend(pm, module.get()); auto &kernelModulePM = pm.nest(); // Inline legal direct calls before VMI layout assignment so private helper // bodies participate in one caller-local layout decision. The Func @@ -3141,8 +2965,6 @@ int mlir::pto::compilePTOASModule( bool emitVPTOHostStub) { result.reset(); std::string arch = resolveEffectiveTargetArch(*module, context.getArch()); - int argc = context.getArgc(); - char **argv = context.getArgv(); // Name-hint provenance: textual .pto inputs had their SSA/arg/block-arg names // attached to op Locations by the driver right after parsing. Collect the @@ -3369,10 +3191,6 @@ int mlir::pto::compilePTOASModule( } const bool hasTileOpsToExpand = hasUnexpandedTileOps(*module); - std::optional expandOptions; - if (effectiveBackend == PTOBackend::VPTO && hasTileOpsToExpand && - tileLibBackend == TileLibBackend::PTODSL) - expandOptions = resolveExpandTileOpOptions(argc, argv); if (effectiveBackend == PTOBackend::VPTO && !hasTileOpsToExpand) { if (ptoPrintSeamIR || !ptoSeamIRFile.empty()) { @@ -3380,8 +3198,7 @@ int mlir::pto::compilePTOASModule( "skipping the shared PTO-to-VPTO lowering pipeline.\n"; return 1; } - if (failed(runVPTOBackendPipeline(module, hasTileOpsToExpand, - /*expandOptions=*/nullptr))) + if (failed(runVPTOBackendPipeline(module, hasTileOpsToExpand))) return 1; return emitVPTOBackendResult(*module, result, emitVPTOHostStub, context.getCANNVersionOrDefault()); @@ -3419,13 +3236,8 @@ int mlir::pto::compilePTOASModule( // PTODSL legality discovery happens on tile-native PTO IR before fusion. // Fusion may later filter the ordered `candidates` array; ExpandTileOp // consumes the first candidate that remains. - if (!isA2A3 && expandOptions && - expandOptions->tileLibBackend == "ptodsl") { - auto insertOptions = - buildInsertTemplateAttributesOptions(*expandOptions); - pm.addPass( - pto::createInsertTemplateAttributesPass(insertOptions)); - } + if (!isA2A3 && effectiveBackend == PTOBackend::VPTO && hasTileOpsToExpand) + pm.addPass(pto::createInsertTemplateAttributesPass()); // Keep frontend fusion on tile-native PTO IR and annotate last_use directly // on scheduled block-local spans before the shared mainline lowers tiles. @@ -3561,12 +3373,6 @@ int mlir::pto::compilePTOASModule( if (ptoPrintSeamIR) printSharedPreBackendSeamIR(*module); - // The PTODSL daemon is needed before the main pipeline for metadata. - // Legacy TileLang can still be resolved lazily immediately before - // ExpandTileOp, preserving the prior --emit-pto-ir behavior. - if (hasTileOpsToExpand && !expandOptions) - expandOptions = resolveExpandTileOpOptions(argc, argv); - if (ptoPrintSeamIR) { module->print(llvm::errs()); llvm::errs() << "\n"; @@ -3574,9 +3380,7 @@ int mlir::pto::compilePTOASModule( if (failed(emitSharedPreBackendSeamIR(*module, ptoSeamIRFile))) return 1; - if (failed(runVPTOBackendPipeline( - module, hasTileOpsToExpand, - expandOptions ? &*expandOptions : nullptr))) + if (failed(runVPTOBackendPipeline(module, hasTileOpsToExpand))) return 1; return emitVPTOBackendResult(*module, result, emitVPTOHostStub, context.getCANNVersionOrDefault()); diff --git a/tools/ptoas/ptoas.h b/tools/ptoas/ptoas.h index cf5b369ef7..500525094c 100644 --- a/tools/ptoas/ptoas.h +++ b/tools/ptoas/ptoas.h @@ -62,6 +62,8 @@ class PTOASContext { public: PTOASContext(DialectRegistry ®istry, llvm::StringRef outputPath, int argc, char **argv); + PTOASContext(MLIRContext &borrowedContext, llvm::StringRef outputPath, + int argc, char **argv); ~PTOASContext(); LogicalResult initializeEnvironment(bool requiresToolchain, @@ -94,7 +96,8 @@ class PTOASContext { std::string &path); private: - MLIRContext mlirContext; + std::unique_ptr ownedMlirContext; + MLIRContext *mlirContext = nullptr; std::string outputPath; std::string arch; BackendInfo backendInfo; @@ -135,6 +138,8 @@ void loadPTOASDialects(MLIRContext &context); // Reusable driver entry shared by the Python extension and standalone CLI. PTOAS_COMPILER_EXPORT int runPTOAS(int argc, char **argv); +PTOAS_COMPILER_EXPORT int +runPTOAS(int argc, char **argv, MLIRContext &borrowedContext); // Attach textual-.pto SSA name hints (function args, block args, op results) // to the parsed module's Locations as debug metadata. Called by the driver