diff --git a/include/PTO/IR/PTOInterfaces.td b/include/PTO/IR/PTOInterfaces.td index ac306203b2..556a13df28 100644 --- a/include/PTO/IR/PTOInterfaces.td +++ b/include/PTO/IR/PTOInterfaces.td @@ -76,6 +76,26 @@ def MteOpInterface : OpInterface<"MteOpInterface"> { operations and raw micro-instructions may both carry this marker. }]; let cppNamespace = "::mlir::pto"; + let methods = [ + InterfaceMethod<"Return the source buffer.", + "::mlir::Value", "getSource", (ins), [{}], + /*defaultImplementation=*/[{ + auto operation = $_op.getOperation(); + if (auto dps = ::llvm::dyn_cast<::mlir::DestinationStyleOpInterface>( + operation)) + return dps.getDpsInputOperand(0)->get(); + return operation->getOperand(0); + }]>, + InterfaceMethod<"Return the destination buffer.", + "::mlir::Value", "getDestination", (ins), [{}], + /*defaultImplementation=*/[{ + auto operation = $_op.getOperation(); + if (auto dps = ::llvm::dyn_cast<::mlir::DestinationStyleOpInterface>( + operation)) + return dps.getDpsInitOperand(0)->get(); + return operation->getOperand(1); + }]> + ]; } def SimtOpInterface : OpInterface<"SimtOpInterface"> { diff --git a/include/PTO/IR/VMIOps.td b/include/PTO/IR/VMIOps.td index b3a2148321..5e3c0e7cdb 100644 --- a/include/PTO/IR/VMIOps.td +++ b/include/PTO/IR/VMIOps.td @@ -40,7 +40,8 @@ def PTO_PhysicalVMIPartTypeConstraint : AnyTypeOf< "PTO physical vector register or mask type">; class VMI_Op traits = []> - : PTO_Op<"vmi." # mnemonic, traits>; + : PTO_Op<"vmi." # mnemonic, + !listconcat(traits, [MicroOpInterface, VectorMicroOpInterface])>; //===--- Legacy (old) VMI ops ---===// diff --git a/include/PTO/IR/VPTOOps.td b/include/PTO/IR/VPTOOps.td index 6b5411e78a..e24371e64f 100644 --- a/include/PTO/IR/VPTOOps.td +++ b/include/PTO/IR/VPTOOps.td @@ -309,7 +309,8 @@ def TileBufAddrOp : PTO_Op<"tile_buf_addr", [Pure]> { }]; } -def VecScopeOp : PTO_Op<"vecscope", [SingleBlock, NoTerminator]> { +def VecScopeOp : PTO_Op<"vecscope", [SingleBlock, NoTerminator, + VectorMicroOpInterface]> { let summary = "Structured region container for one VPTO vector scope"; let description = [{ `pto.vecscope` marks a structured vector-scope interval without overloading @@ -324,7 +325,8 @@ def VecScopeOp : PTO_Op<"vecscope", [SingleBlock, NoTerminator]> { } def StrictVecScopeOp : PTO_Op<"strict_vecscope", [SingleBlock, NoTerminator, - IsolatedFromAbove]> { + IsolatedFromAbove, + VectorMicroOpInterface]> { let summary = "Structured VPTO vector scope with explicit captures only"; let description = [{ `pto.strict_vecscope` is the strict form of `pto.vecscope`. Values used by diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index 29a43cee84..4ebfe65ec5 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -355,7 +355,7 @@ def PTONormalizeUncoveredTileSections This pass is the shared entry point for inferred-section normalization: explicit sections remain in place, while uncovered top-level segments that - contain tile operations are identified for later wrapping. + contain section-relevant operations are identified for later wrapping. }]; let constructor = "mlir::pto::createPTONormalizeUncoveredTileSectionsPass()"; diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index af6dd28592..36605c6d13 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -3589,11 +3589,15 @@ LogicalResult mlir::pto::SyncSetOp::verify() { } switch (getPipe().getPipe()) { case PIPE::PIPE_FIX: + case PIPE::PIPE_MTE1: + case PIPE::PIPE_MTE2: case PIPE::PIPE_MTE3: + case PIPE::PIPE_V: return success(); default: - return emitOpError() - << "A5 sync.set expects pipe to be one of , "; + return emitOpError() << "A5 sync.set expects pipe to be one of " + ", , , " + ", "; } }; return dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5); diff --git a/lib/PTO/Transforms/PTONormalizeUncoveredTileSections.cpp b/lib/PTO/Transforms/PTONormalizeUncoveredTileSections.cpp index c9a5acf953..5dda7b7aaa 100644 --- a/lib/PTO/Transforms/PTONormalizeUncoveredTileSections.cpp +++ b/lib/PTO/Transforms/PTONormalizeUncoveredTileSections.cpp @@ -1,10 +1,12 @@ // 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. +// 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/IR/PTO.h" #include "PTO/Transforms/Passes.h" @@ -69,6 +71,100 @@ static bool isTileLikeOp(Operation *op) { op->getName().getStringRef().starts_with("pto.t"); } +// Low-level PTODSL sources do not contain pto.t* TileOps. Keep their +// section inference conservative: only operations with an unambiguous engine +// ownership are treated as section carriers. +static bool isRawSectionCarrierOp(Operation *op) { + return op && + isa(op); +} + +static std::optional +classifyTileOpByAddressSpace(Operation *op); +static std::optional classifyTileOpByPipe(Operation *op); +static std::optional getBufferAddressSpace(Type type); + +static std::optional classifySyncPipe(PIPE pipe) { + switch (pipe) { + case PIPE::PIPE_M: + case PIPE::PIPE_MTE1: + case PIPE::PIPE_FIX: + return InferredSectionKind::Cube; + case PIPE::PIPE_V: + case PIPE::PIPE_V2: + case PIPE::PIPE_MTE3: + return InferredSectionKind::Vector; + // Both cores have scalar and MTE2 pipelines. Their physical section must + // come from the peer synchronization pipe or the operation's buffers. + case PIPE::PIPE_S: + case PIPE::PIPE_MTE2: + default: + return std::nullopt; + } +} + +static std::optional +classifyFlagPipes(PipeAttr source, PipeAttr destination) { + std::optional sourceKind = + classifySyncPipe(source.getPipe()); + std::optional destinationKind = + classifySyncPipe(destination.getPipe()); + if (sourceKind && destinationKind && *sourceKind != *destinationKind) + return std::nullopt; + return sourceKind ? sourceKind : destinationKind; +} + +static std::optional +classifyMteOpByAddressSpace(MteOpInterface mteOp) { + std::optional source = + getBufferAddressSpace(mteOp.getSource().getType()); + std::optional destination = + getBufferAddressSpace(mteOp.getDestination().getType()); + if (!source || !destination) + return std::nullopt; + + if (*source == AddressSpace::ACC) + return InferredSectionKind::Cube; + if (*source == AddressSpace::VEC || *destination == AddressSpace::VEC) + return InferredSectionKind::Vector; + if (*destination == AddressSpace::MAT || *destination == AddressSpace::LEFT || + *destination == AddressSpace::RIGHT || + *destination == AddressSpace::BIAS || + *destination == AddressSpace::SCALING) + return InferredSectionKind::Cube; + return std::nullopt; +} + +static std::optional +classifyRawSectionCarrierOp(Operation *op) { + if (!isRawSectionCarrierOp(op)) + return std::nullopt; + if (isa(op)) + return InferredSectionKind::Vector; + if (isa(op)) + return InferredSectionKind::Cube; + if (auto mteOp = dyn_cast(op)) { + if (auto kind = classifyMteOpByAddressSpace(mteOp)) + return kind; + return classifyTileOpByPipe(op); + } + if (auto setFlag = dyn_cast(op)) + return classifyFlagPipes(setFlag.getSrcPipe(), setFlag.getDstPipe()); + if (auto waitFlag = dyn_cast(op)) + return classifyFlagPipes(waitFlag.getSrcPipe(), waitFlag.getDstPipe()); + if (auto setFlag = dyn_cast(op)) + return classifyFlagPipes(setFlag.getSrcPipe(), setFlag.getDstPipe()); + if (auto waitFlag = dyn_cast(op)) + return classifyFlagPipes(waitFlag.getSrcPipe(), waitFlag.getDstPipe()); + if (auto syncSet = dyn_cast(op)) + return classifySyncPipe(syncSet.getPipe().getPipe()); + if (auto syncWait = dyn_cast(op)) + return classifySyncPipe(syncWait.getPipe().getPipe()); + return std::nullopt; +} + static bool isPipeLikeOp(Operation *op) { return op && isa(op); } @@ -104,8 +200,8 @@ static bool hasAnySection(func::FuncOp funcOp) { } static bool hasExplicitFunctionKernelKind(func::FuncOp funcOp) { - return funcOp && - funcOp->hasAttrOfType(FunctionKernelKindAttr::name); + return funcOp && funcOp->hasAttrOfType( + FunctionKernelKindAttr::name); } static bool isInsideKernelKindModule(func::FuncOp funcOp) { @@ -116,10 +212,13 @@ static bool isInsideKernelKindModule(func::FuncOp funcOp) { } static bool hasKnownKernelKindContext(func::FuncOp funcOp) { - return isInsideKernelKindModule(funcOp) || hasExplicitFunctionKernelKind(funcOp); + return isInsideKernelKindModule(funcOp) || + hasExplicitFunctionKernelKind(funcOp); } static std::optional getBufferAddressSpace(Type type) { + if (auto ptrType = dyn_cast(type)) + return ptrType.getMemorySpace().getAddressSpace(); if (auto tileType = dyn_cast(type)) { if (auto attr = dyn_cast_or_null(tileType.getMemorySpace())) @@ -240,8 +339,7 @@ static std::optional classifyTileOpByName(Operation *op) { return std::nullopt; } -static std::optional -classifyTileOpByPipe(Operation *op) { +static std::optional classifyTileOpByPipe(Operation *op) { return inferPhysicalSectionKindFromPipe(op); } @@ -348,7 +446,8 @@ static std::optional classifyTileOp(Operation *op) { if (std::optional kind = classifyTStoreBySourceAddressSpace(op)) return kind; - if (std::optional kind = classifyTileOpByAddressSpace(op)) + if (std::optional kind = + classifyTileOpByAddressSpace(op)) return kind; return classifyTileOpByPipe(op); } @@ -366,7 +465,8 @@ enum class FunctionKindCacheState : uint8_t { InProgress = 3, }; -static void inspectModuleKindOperation(Operation *op, ModuleKindSummary &summary) { +static void inspectModuleKindOperation(Operation *op, + ModuleKindSummary &summary) { if (!op) return; if (isExplicitSection(op)) @@ -446,9 +546,9 @@ static func::CallOp getTransparentWrapperCall(func::FuncOp funcOp) { return callOp; } -static std::optional -inferWholeFunctionKind(func::FuncOp funcOp, - llvm::DenseMap &cache) { +static std::optional inferWholeFunctionKind( + func::FuncOp funcOp, + llvm::DenseMap &cache) { if (!funcOp || funcOp.isDeclaration()) return std::nullopt; @@ -463,7 +563,8 @@ inferWholeFunctionKind(func::FuncOp funcOp, ModuleKindSummary summary; inspectModuleKindOperation(funcOp.getOperation(), summary); std::optional inferredKind; - if (summary.ambiguousOps.empty() && !(summary.vectorCount && summary.cubeCount)) { + if (summary.ambiguousOps.empty() && + !(summary.vectorCount && summary.cubeCount)) { if (summary.vectorCount) inferredKind = InferredSectionKind::Vector; else if (summary.cubeCount) @@ -472,8 +573,8 @@ inferWholeFunctionKind(func::FuncOp funcOp, if (!inferredKind) { if (func::CallOp callOp = getTransparentWrapperCall(funcOp)) { - auto callee = - SymbolTable::lookupNearestSymbolFrom(funcOp, callOp.getCalleeAttr()); + auto callee = SymbolTable::lookupNearestSymbolFrom( + funcOp, callOp.getCalleeAttr()); if (callee && callee != funcOp) inferredKind = inferWholeFunctionKind(callee, cache); } @@ -484,9 +585,9 @@ inferWholeFunctionKind(func::FuncOp funcOp, } static void assignModuleKernelKind(ModuleOp module, InferredSectionKind kind) { - FunctionKernelKind kernelKind = - kind == InferredSectionKind::Vector ? FunctionKernelKind::Vector - : FunctionKernelKind::Cube; + FunctionKernelKind kernelKind = kind == InferredSectionKind::Vector + ? FunctionKernelKind::Vector + : FunctionKernelKind::Cube; module->setAttr(FunctionKernelKindAttr::name, FunctionKernelKindAttr::get(module.getContext(), kernelKind)); } @@ -496,9 +597,9 @@ static void assignFunctionKernelKind(func::FuncOp funcOp, if (!funcOp) return; - FunctionKernelKind kernelKind = - kind == InferredSectionKind::Vector ? FunctionKernelKind::Vector - : FunctionKernelKind::Cube; + FunctionKernelKind kernelKind = kind == InferredSectionKind::Vector + ? FunctionKernelKind::Vector + : FunctionKernelKind::Cube; funcOp->setAttr(FunctionKernelKindAttr::name, FunctionKernelKindAttr::get(funcOp.getContext(), kernelKind)); } @@ -544,7 +645,8 @@ static LogicalResult tryAssignWholeFunctionKernelKind(func::FuncOp funcOp) { return success(); llvm::DenseMap cache; - std::optional kind = inferWholeFunctionKind(funcOp, cache); + std::optional kind = + inferWholeFunctionKind(funcOp, cache); if (!kind) return success(); @@ -557,9 +659,11 @@ static void inspectSegmentOperation(Operation *op, if (!op) return; - if (isTileLikeOp(op)) { + if (isTileLikeOp(op) || isRawSectionCarrierOp(op)) { segment.containsTileOp = true; - if (std::optional kind = classifyTileOp(op)) { + std::optional kind = + isTileLikeOp(op) ? classifyTileOp(op) : classifyRawSectionCarrierOp(op); + if (kind) { if (*kind == InferredSectionKind::Vector) ++segment.vectorTileOpCount; else @@ -663,8 +767,9 @@ static void collectUncoveredTopLevelSegments( } template -static void wrapUncoveredTopLevelSegment(func::FuncOp funcOp, - const UncoveredTopLevelSegment &segment) { +static void +wrapUncoveredTopLevelSegment(func::FuncOp funcOp, + const UncoveredTopLevelSegment &segment) { Block &entryBlock = funcOp.getBody().front(); Operation *firstOp = segment.firstOp; Operation *lastOp = segment.lastOp; @@ -678,40 +783,44 @@ static void wrapUncoveredTopLevelSegment(func::FuncOp funcOp, auto firstIt = Block::iterator(firstOp); auto afterLastIt = std::next(Block::iterator(lastOp)); - sectionBlock->getOperations().splice(sectionBlock->end(), - entryBlock.getOperations(), firstIt, - afterLastIt); + sectionBlock->getOperations().splice( + sectionBlock->end(), entryBlock.getOperations(), firstIt, afterLastIt); } -static LogicalResult emitSegmentInferenceError(func::FuncOp funcOp, - const UncoveredTopLevelSegment &segment) { +static LogicalResult +emitSegmentInferenceError(func::FuncOp funcOp, + const UncoveredTopLevelSegment &segment) { InFlightDiagnostic diag = - funcOp.emitOpError("contains an uncovered top-level TileOp segment whose " + funcOp.emitOpError("contains an uncovered top-level op segment whose " "section kind cannot be inferred uniquely"); if (segment.vectorTileOpCount && segment.cubeTileOpCount) { - diag << "; saw both vector-like and cube-like TileOps in the same segment"; + diag << "; saw both vector-like and cube-like ops in the same segment"; } else if (!segment.ambiguousTileOps.empty()) { - diag << "; ambiguous TileOp(s): "; - for (size_t i = 0, e = segment.ambiguousTileOps.size(); i < e && i < 3; ++i) { + diag << "; ambiguous op(s): "; + for (size_t i = 0, e = segment.ambiguousTileOps.size(); i < e && i < 3; + ++i) { if (i) diag << ", "; diag << '\'' << segment.ambiguousTileOps[i]->getName().getStringRef() << '\''; } } + diag << "; wrap the ambiguous region in pto.section.cube or " + "pto.section.vector to specify its physical section explicitly"; return failure(); } -static LogicalResult emitResidualUncoveredTileSegmentError( - func::FuncOp funcOp, const UncoveredTopLevelSegment &segment) { +static LogicalResult +emitResidualUncoveredTileSegmentError(func::FuncOp funcOp, + const UncoveredTopLevelSegment &segment) { InFlightDiagnostic diag = funcOp.emitOpError( - "still contains an uncovered top-level TileOp segment after section " + "still contains an uncovered top-level op segment after section " "normalization"); if (segment.containsNestedExplicitSection) { diag << "; a top-level op mixes nested explicit pto.section.* with sibling " - "TileOps outside those sections"; + "ops outside those sections"; } - diag << "; first uncovered TileOp segment starts at '" + diag << "; first uncovered op segment starts at '" << (segment.firstTileCarrierOp ? segment.firstTileCarrierOp : segment.firstOp) ->getName() @@ -746,8 +855,8 @@ static LogicalResult normalizeFunction(func::FuncOp funcOp) { return success(); } -static LogicalResult verifyFunctionHasNoResidualUncoveredTileSegments( - func::FuncOp funcOp) { +static LogicalResult +verifyFunctionHasNoResidualUncoveredTileSegments(func::FuncOp funcOp) { if (hasKnownKernelKindContext(funcOp)) return success(); diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 33e5e66439..3ff9ed18d7 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -6479,7 +6479,7 @@ struct PTOSyncSetToEmitC : public OpConversionPattern { return success(); } - [[maybe_unused]] PTOArch targetArch; + PTOArch targetArch; }; struct PTOSyncWaitToEmitC : public OpConversionPattern { @@ -6514,7 +6514,7 @@ struct PTOSyncWaitToEmitC : public OpConversionPattern { return success(); } - [[maybe_unused]] PTOArch targetArch; + PTOArch targetArch; }; // GetBlockIdxOp Lowering (pto.get_block_idx -> get_block_idx()) diff --git a/lib/PTO/Transforms/VMILayoutSupport.cpp b/lib/PTO/Transforms/VMILayoutSupport.cpp index 4e5dcb6c35..97d1da0142 100644 --- a/lib/PTO/Transforms/VMILayoutSupport.cpp +++ b/lib/PTO/Transforms/VMILayoutSupport.cpp @@ -311,7 +311,7 @@ static constexpr EnsureLayoutPattern kEnsureLayoutPatterns[] = { {bits<8>(), N<1, 2, 4, 8, 64, 128>(), c(), ls(2)}, {bits<16>(), N<1, 2, 4, 8, 64>(), c(), ls(2)}, {bits<8>(), N<1, 2, 4, 8, 64, 128>(), ls(2), c()}, - {bits<16>(), N<1, 2, 4, 8, 64>(), ls(2), c()}, + {bits<16>(), N<1, 2, 4, 8, 64, 128, 256>(), ls(2), c()}, {bits<8>(), N<1, 2, 4, 8, 64>(), c(), ls(4)}, {bits<8>(), N<1, 2, 4, 8, 64>(), ls(4), c()}, diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index b76bbd3251..73324e4602 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -3620,10 +3620,11 @@ FailureOr> materializeContiguousToLaneStride( FailureOr> materializeLaneStrideToContiguous( Operation *op, ValueRange sourceParts, TypeRange resultTypes, Type elementType, int64_t laneStride, PatternRewriter &rewriter) { - if (sourceParts.size() != resultTypes.size()) { + if (sourceParts.empty() || resultTypes.empty() || + sourceParts.size() > resultTypes.size() * laneStride) { (void)rewriter.notifyMatchFailure( - op, "dense lane_stride pack materialization requires matching " - "source/result physical arity"); + op, "dense lane_stride pack materialization source arity does not " + "fit result arity"); return failure(); } @@ -3644,26 +3645,58 @@ FailureOr> materializeLaneStrideToContiguous( return failure(); SmallVector results; - results.reserve(sourceParts.size()); - for (auto [source, resultType] : llvm::zip_equal(sourceParts, resultTypes)) { - FailureOr current = - bitcastVReg(op->getLoc(), source, *sourceCarrier, rewriter); - if (failed(current)) + results.reserve(resultTypes.size()); + auto packPart = [&](Value source, StringRef part, + unsigned inputBits) -> FailureOr { + FailureOr resultCarrier = + getUnsignedCarrierVRegType(rewriter.getContext(), inputBits / 2); + if (failed(resultCarrier)) return failure(); - FailureOr packed = packToPreviousCarrier(op->getLoc(), *current, - carrierBits / 2, rewriter); - if (failed(packed)) + return rewriter + .create(op->getLoc(), *resultCarrier, source, + rewriter.getStringAttr(part)) + .getResult(); + }; + for (auto [resultIndex, resultType] : llvm::enumerate(resultTypes)) { + size_t sourceBegin = resultIndex * laneStride; + if (sourceBegin >= sourceParts.size()) return failure(); - current = *packed; + size_t sourceEnd = std::min(sourceBegin + static_cast(laneStride), + sourceParts.size()); + Value current; + for (size_t sourceIndex = sourceBegin; sourceIndex < sourceEnd; + ++sourceIndex) { + FailureOr source = bitcastVReg( + op->getLoc(), sourceParts[sourceIndex], *sourceCarrier, rewriter); + if (failed(source)) + return failure(); + FailureOr packed = packPart( + *source, sourceIndex == sourceBegin ? "LOWER" : "HIGHER", + carrierBits); + if (failed(packed)) + return failure(); + if (!current) { + current = *packed; + } else { + auto packedType = cast(packed->getType()); + FailureOr mask = + createAllTrueMaskForVReg(op->getLoc(), packedType, rewriter); + if (failed(mask)) + return failure(); + current = rewriter + .create(op->getLoc(), packedType, current, *packed, + *mask) + .getResult(); + } + } if (laneStride == 4) { - packed = - packToPreviousCarrier(op->getLoc(), *current, elementBits, rewriter); + FailureOr packed = packPart(current, "LOWER", elementBits * 2); if (failed(packed)) return failure(); current = *packed; } FailureOr result = - bitcastVReg(op->getLoc(), *current, resultType, rewriter); + bitcastVReg(op->getLoc(), current, resultType, rewriter); if (failed(result)) return failure(); results.push_back(*result); diff --git a/lib/PTO/Transforms/VPTOSplitCVModule.cpp b/lib/PTO/Transforms/VPTOSplitCVModule.cpp index fcea35acd6..b57eb4f215 100644 --- a/lib/PTO/Transforms/VPTOSplitCVModule.cpp +++ b/lib/PTO/Transforms/VPTOSplitCVModule.cpp @@ -1,10 +1,12 @@ // 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. +// 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/IR/PTO.h" #include "PTO/Transforms/Passes.h" @@ -36,6 +38,31 @@ static bool hasKernelKindChildModule(ModuleOp module) { [](ModuleOp child) { return hasKernelKind(child); }); } +static bool hasCVSections(ModuleOp module); + +static bool flattenSingleUnpartitionedChild(ModuleOp module) { + SmallVector topLevelOps; + for (Operation &op : module.getBodyRegion().front().getOperations()) + topLevelOps.push_back(&op); + if (topLevelOps.size() != 1) + return false; + + auto child = dyn_cast(topLevelOps.front()); + if (!child || hasKernelKind(child) || !hasCVSections(child)) + return false; + + SmallVector childAttrs(child->getAttrs().begin(), + child->getAttrs().end()); + Region childBody; + childBody.takeBody(child.getBodyRegion()); + child.erase(); + module.getBodyRegion().takeBody(childBody); + for (NamedAttribute attr : childAttrs) + if (attr.getName() != SymbolTable::getSymbolAttrName()) + module->setAttr(attr.getName(), attr.getValue()); + return true; +} + static bool isSectionSplitCandidate(func::FuncOp funcOp); static bool hasCVSections(ModuleOp module) { @@ -163,7 +190,8 @@ static LogicalResult verifyExplicitKernelKindMatchesSections(ModuleOp module) { if (!isCube && !isVector) return WalkResult::advance(); if (isCube != expectsCube) { - status = op->emitError("conflicts with explicit pto.kernel_kind on its module"); + status = op->emitError( + "conflicts with explicit pto.kernel_kind on its module"); return WalkResult::interrupt(); } return WalkResult::advance(); @@ -187,8 +215,9 @@ static LogicalResult verifySectionSplitCandidatesUseSections(ModuleOp module) { return status; } -static void eraseSectionSplitCandidatesWithoutSectionKind(ModuleOp module, - FunctionKernelKind kind) { +static void +eraseSectionSplitCandidatesWithoutSectionKind(ModuleOp module, + FunctionKernelKind kind) { SmallVector eraseFuncs; module.walk([&](func::FuncOp funcOp) { if (isSectionSplitCandidate(funcOp) && !hasSectionKind(funcOp, kind)) @@ -257,6 +286,7 @@ static LogicalResult materializeExplicitKernelKindSections(ModuleOp module) { } static LogicalResult splitCVModule(ModuleOp module) { + flattenSingleUnpartitionedChild(module); if (hasKernelKind(module)) return materializeExplicitKernelKindSections(module); if (hasKernelKindChildModule(module)) { diff --git a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md index 604ee52365..3e96a6ea2e 100644 --- a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md +++ b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md @@ -652,6 +652,29 @@ you can drop below the tile abstraction without leaving the `@pto.jit` entry. When you want to delineate an authored vector interval directly in the public PTODSL surface, use `with pto.vecscope():` inside an explicit-mode kernel body. +PTOAS normally infers whether uncovered Tile and micro-instruction regions run +on the Cube or Vector core. If a region only contains operations shared by both +cores, such as synchronization between `PIPE_MTE2` and `PIPE_S`, inference +cannot determine the physical core uniquely. Use `with pto.section("cube"):` +or `with pto.section("vector"):` around that complete region as an explicit +placement hint: + +```python +@pto.jit(target="a5", mode="explicit") +def mixed_kernel(): + with pto.section("cube"): + pto.set_flag("MTE2", "S", event_id=0) + + with pto.section("vector"): + pto.wait_flag("MTE2", "S", event_id=0) +``` + +Explicit sections are an escape hatch rather than a requirement: regions that +PTOAS can classify should remain uncovered and continue to use automatic +inference. Sections cannot be nested, each kind may appear at most once per +function, and `pto.section()` cannot be combined with an explicit +`@pto.jit(kernel_kind=...)` contract. + The richer type surface also applies to sub-kernels: in auto mode, a sub-kernel's parameters are restricted to `Tile` and PTO scalar types; in explicit mode they may also accept `PartitionTensorView` and `pto.ptr`, diff --git a/ptodsl/docs/user_guide/10-sync-ops.md b/ptodsl/docs/user_guide/10-sync-ops.md index 1d7af845de..dcd6917ef8 100644 --- a/ptodsl/docs/user_guide/10-sync-ops.md +++ b/ptodsl/docs/user_guide/10-sync-ops.md @@ -403,7 +403,7 @@ pto.wait_cross_flag(pto.Pipe.FIX, 0) ### 10.5.2 Intra-block sync: `set_intra_flag`, `wait_intra_flag` -The Cube unit (matrix pipeline) has a dedicated synchronization channel separate from the standard pipe-flag mechanism used by MTE and Vector pipelines. `set_intra_flag` and `wait_intra_flag` synchronize Cube and Vector within the same block, ensuring that shared UB tile data is not accessed before the producer finishes. +The intra-block sync channel is separate from the standard pipe-flag mechanism used by cross-core sync. `set_intra_flag` and `wait_intra_flag` synchronize the relevant producer/consumer pipes within the same block, ensuring that shared UB tile data is not accessed before the producer finishes. Cross-core flags use the public `0`-`7` event range. A5 intra-block sync uses a separate physical event range, described below. @@ -418,7 +418,7 @@ Unlike `wait_cross_flag`, `wait_intra_flag` only stalls the specified pipeline | Parameter | Type | Description | |-----------|------|-------------| -| `pipe` | `Pipe` | Trigger endpoint for the synchronization event. The public DSL accepts `Pipe.FIX` and `Pipe.MTE3` here. | +| `pipe` | `Pipe` | Trigger endpoint for the synchronization event. The public DSL accepts `Pipe.FIX`, `Pipe.MTE1`, `Pipe.MTE2`, `Pipe.MTE3`, and `Pipe.V` here. | | `event_id` | `int` or scalar expression | Physical event identifier (`0`–`31` for static IDs). | For A5 mixed C/V writeback code that must notify both AIV subblocks, emit two signals @@ -442,7 +442,7 @@ pto.set_intra_flag(pto.Pipe.MTE3, 0) | Parameter | Type | Description | |-----------|------|-------------| -| `pipe` | `Pipe` | Waiting endpoint for the synchronization event. The public DSL accepts `Pipe.FIX`, `Pipe.MTE3`, and `Pipe.V` here. | +| `pipe` | `Pipe` | Waiting endpoint for the synchronization event. The public DSL accepts `Pipe.FIX`, `Pipe.MTE1`, `Pipe.MTE2`, `Pipe.MTE3`, and `Pipe.V` here. | | `event_id` | `int` or scalar expression | Physical event identifier to wait on (`0`–`31` for static IDs). | For A5 mixed C/V writeback code that must synchronize with both AIV subblocks, wait on both diff --git a/ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md b/ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md index 6229bb953f..b11ee24550 100644 --- a/ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md +++ b/ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md @@ -805,8 +805,8 @@ group_max = pto.vmi.vcmax( - `reassoc` is only meaningful for `vcadd` on floating-point data. - Floating-point `vcadd` must spell `reassoc` explicitly at the PTODSL surface. - `reassoc=None` is rejected by PTODSL; use `reassoc=True` or `reassoc=False`. -- The current VMI op encoding remains presence-based, so `reassoc=False` - lowers to the same no-attribute form as legacy callers. +- The current VMI op encoding remains presence-based, so PTODSL preserves the + `reassoc` attribute for both `reassoc=True` and `reassoc=False`. --- diff --git a/ptodsl/ptoas_wheel_bootstrap.py b/ptodsl/ptoas_wheel_bootstrap.py index d9da235a63..e7616fbba7 100644 --- a/ptodsl/ptoas_wheel_bootstrap.py +++ b/ptodsl/ptoas_wheel_bootstrap.py @@ -67,7 +67,12 @@ def _candidate_wheel_python_roots(wrapper: Path, module_file: str | None = None) candidates: list[Path] = [] if module_file: - candidates.append(Path(module_file).resolve().parent) + module_path = Path(module_file).resolve() + candidates.append(module_path.parent) + # Editable installs import this bootstrap from the repository while + # the generated runtime overlay lives under build/python. + repo_root = module_path.parent.parent + candidates.append(repo_root / "build" / "python") for scheme_name in ("purelib", "platlib"): scheme_path = sysconfig.get_path(scheme_name) diff --git a/ptodsl/ptodsl/_control_flow.py b/ptodsl/ptodsl/_control_flow.py index 55f389840c..8f30c7932b 100644 --- a/ptodsl/ptodsl/_control_flow.py +++ b/ptodsl/ptodsl/_control_flow.py @@ -12,6 +12,7 @@ Public API ────────── +``section(kind)`` – ``pto.section.cube/vector { … }`` ``vecscope()`` – ``pto.vecscope { … }`` ``for_(lo, hi, step)`` – ``scf.for`` with optional named carry state via ``.carry(...)`` @@ -26,7 +27,7 @@ from ._runtime_index_ops import coerce_runtime_index from ._scalar_coercion import coerce_scalar_to_type from ._surface_types import const_expr -from ._tracing.active import current_session +from ._tracing.active import current_session, require_active_session from ._surface_values import unwrap_surface_value, wrap_like_surface_value, wrap_surface_value from mlir.dialects import pto as _pto, scf @@ -63,6 +64,33 @@ def vecscope() -> _VecScopeCM: return _VecScopeCM() +class _SectionCM: + """Context manager for an explicit cube or vector physical section.""" + + def __init__(self, kind: str): + if not isinstance(kind, str): + raise TypeError("pto.section(kind) expects 'cube' or 'vector'") + if kind not in {"cube", "vector"}: + raise ValueError("pto.section(kind) expects 'cube' or 'vector'") + self._kind = kind + self._cm = None + + def __enter__(self): + _require_explicit_mode("pto.section()") + session = require_active_session("pto.section") + self._cm = session.enter_physical_section(self._kind) + self._cm.__enter__() + return None + + def __exit__(self, *exc): + return self._cm.__exit__(*exc) + + +def section(kind: str) -> _SectionCM: + """Explicitly assign the enclosed operations to the cube or vector core.""" + return _SectionCM(kind) + + # ── compile-time control-flow helpers ───────────────────────────────────────── def static_range(*args): @@ -638,6 +666,6 @@ def yield_(*vals): __all__ = [ - "vecscope", "static_range", "const_expr", "LoopHandle", "BranchHandle", + "section", "vecscope", "static_range", "const_expr", "LoopHandle", "BranchHandle", "for_", "if_", "yield_", ] diff --git a/ptodsl/ptodsl/_jit.py b/ptodsl/ptodsl/_jit.py index 26fb8311e5..41f19a6eb0 100644 --- a/ptodsl/ptodsl/_jit.py +++ b/ptodsl/ptodsl/_jit.py @@ -31,10 +31,10 @@ from mlir.ir import InsertionPoint -_MODULE_ATTRS = ("pto.target_arch",) +_MODULE_ATTRS = ("pto.target_arch", "pto.backend") _SUPPORTED_FRONTEND_OPTION_KEYS = {"ast_rewrite", "rewrite_part", "dump_rewritten_source"} _SUPPORTED_REWRITE_PARTS = {"control_flow"} -_DEFAULT_KERNEL_KIND = "vector" +_DEFAULT_KERNEL_KIND = None class _DefaultKernelKindSentinel: @@ -139,14 +139,22 @@ def merge_jit_modules(*kernels: KernelHandle): """ Merge multiple ``@pto.jit`` kernels into one MLIR module. - Each handle must have been compiled with compatible outer-container - attributes. Child module order follows *kernels*. + Each handle must have compatible top-level attributes and use the same + flat or backend-partitioned module shape. Payload order follows *kernels*. """ if not kernels: raise ValueError("merge_jit_modules() requires at least one kernel handle") merged = kernels[0].build() expected_attrs = _module_attr_map(merged) + merged_is_partitioned = all( + op.operation.name == "builtin.module" for op in merged.body.operations + ) + merged_symbols = { + str(op.operation.attributes["sym_name"]) + for op in merged.body.operations + if "sym_name" in op.operation.attributes + } for kernel in kernels[1:]: module = kernel.build() @@ -156,13 +164,28 @@ def merge_jit_modules(*kernels: KernelHandle): "merge_jit_modules() requires compatible module attributes; " f"expected {expected_attrs}, got {actual_attrs}" ) + module_is_partitioned = all( + op.operation.name == "builtin.module" for op in module.body.operations + ) + if module_is_partitioned != merged_is_partitioned: + raise ValueError( + "merge_jit_modules() cannot mix flat and backend-partitioned PTODSL modules" + ) + if not merged_is_partitioned: + module_symbols = { + str(op.operation.attributes["sym_name"]) + for op in module.body.operations + if "sym_name" in op.operation.attributes + } + duplicates = merged_symbols & module_symbols + if duplicates: + raise ValueError( + "merge_jit_modules() cannot merge duplicate symbols into a flat module: " + f"{sorted(duplicates)!r}" + ) + merged_symbols.update(module_symbols) with InsertionPoint(merged.body): for op in module.body.operations: - if op.operation.name != "builtin.module": - raise ValueError( - "merge_jit_modules() expects backend-partitioned PTODSL containers " - "whose payload is expressed entirely through child modules" - ) op.operation.clone() merged.operation.verify() @@ -173,7 +196,7 @@ def jit( name=None, *, target: str = "a5", - kernel_kind: str = _DEFAULT_KERNEL_KIND_SENTINEL, + kernel_kind: str | None = _DEFAULT_KERNEL_KIND_SENTINEL, backend: str = "vpto", entry: bool = True, mode: str = "auto", @@ -191,8 +214,8 @@ def jit( target: Target architecture string, e.g. ``"a5"``. kernel_kind: optional authored physical kind, used for native build selection and explicit single-kind VPTO authoring intent. When omitted, - PTODSL keeps the historical vector default while allowing - subkernel sections to express mixed cube/vector regions. + PTODSL leaves the kind unspecified so PTOAS can split mixed + cube/vector regions. backend: ``"vpto"`` or ``"emitc"`` – records the intended backend. entry: ``True`` for launchable kernel entries, ``False`` for helpers. mode: ``"auto"`` or ``"explicit"`` – feeds child compile policy. diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index c3349f5459..7e1ec04ad1 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -6136,7 +6136,7 @@ def set_intra_flag(pipe, event_id): _validate_sync_pipe( pipe, context="set_intra_flag(pipe, event_id)", - allowed=("PIPE_FIX", "PIPE_MTE3"), + allowed=("PIPE_FIX", "PIPE_MTE1", "PIPE_MTE2", "PIPE_MTE3", "PIPE_V"), ) event_operand = _sync_event_id_operand_in_range( event_id, @@ -6153,7 +6153,7 @@ def wait_intra_flag(pipe, event_id): _validate_sync_pipe( pipe, context="wait_intra_flag(pipe, event_id)", - allowed=("PIPE_FIX", "PIPE_MTE3", "PIPE_V"), + allowed=("PIPE_FIX", "PIPE_MTE1", "PIPE_MTE2", "PIPE_MTE3", "PIPE_V"), ) event_operand = _sync_event_id_operand_in_range( event_id, diff --git a/ptodsl/ptodsl/_runtime/native_build.py b/ptodsl/ptodsl/_runtime/native_build.py index 310b44067e..22516c7eb7 100644 --- a/ptodsl/ptodsl/_runtime/native_build.py +++ b/ptodsl/ptodsl/_runtime/native_build.py @@ -122,7 +122,7 @@ def _host_compile_flags() -> list[str]: ] -def _kernel_compile_flags(kernel_kind: str, target_arch: str) -> list[str]: +def _kernel_compile_flags(kernel_kind: str | None, target_arch: str) -> list[str]: arch = aicore_arch_for_kernel_kind(kernel_kind, target_arch) return common_include_flags() + [ "-std=gnu++17", @@ -152,7 +152,7 @@ def _compile_launch_cpp( launch_cpp: Path, launch_object: Path, *, - kernel_kind: str, + kernel_kind: str | None, target_arch: str, export_macro: str, ) -> None: @@ -175,7 +175,7 @@ def _link_shared_library( kernel_object: Path, shared_library: Path, *, - kernel_kind: str, + kernel_kind: str | None, ) -> None: bisheng = resolve_bisheng() soname = shared_library.name diff --git a/ptodsl/ptodsl/_runtime/toolchain.py b/ptodsl/ptodsl/_runtime/toolchain.py index b41ef53a54..b40568a074 100644 --- a/ptodsl/ptodsl/_runtime/toolchain.py +++ b/ptodsl/ptodsl/_runtime/toolchain.py @@ -170,8 +170,10 @@ def runtime_library_flags(*, sim_mode: bool = False) -> list[str]: return flags -def aicore_arch_for_kernel_kind(kernel_kind: str, target_arch: str) -> str: +def aicore_arch_for_kernel_kind(kernel_kind: str | None, target_arch: str) -> str: target = target_arch.lower() + if kernel_kind is None: + return "dav-c220" if target in {"a2", "a3"} else "dav-c310" if kernel_kind == "vector": if target in {"a2", "a3"}: return "dav-c220-vec" diff --git a/ptodsl/ptodsl/_tracing/module_builder.py b/ptodsl/ptodsl/_tracing/module_builder.py index e5fb99859a..15f48ee8f4 100644 --- a/ptodsl/ptodsl/_tracing/module_builder.py +++ b/ptodsl/ptodsl/_tracing/module_builder.py @@ -30,7 +30,7 @@ class KernelModuleSpec: function_name: str target_arch: str - kernel_kind: str + kernel_kind: str | None kernel_kind_explicit: bool = False backend: str = "vpto" entry: bool = True @@ -97,6 +97,15 @@ def _build_backend_partitioned_module(spec: KernelModuleSpec, arg_types): outer = Module.create() outer.operation.attributes["pto.target_arch"] = StringAttr.get(spec.target_arch) + if spec.kernel_kind is None: + outer.operation.attributes["pto.backend"] = StringAttr.get(spec.backend) + fn_ty = func.FunctionType.get(arg_types, []) + with InsertionPoint(outer.body): + ir_fn = func.FuncOp(spec.function_name, fn_ty) + if spec.entry: + ir_fn.attributes["pto.entry"] = UnitAttr.get() + return outer, ir_fn + _, child_body = create_container_child_module(outer, spec) fn_ty = func.FunctionType.get(arg_types, []) with InsertionPoint(child_body): diff --git a/ptodsl/ptodsl/_tracing/session.py b/ptodsl/ptodsl/_tracing/session.py index 617fc05578..4e9b29db43 100644 --- a/ptodsl/ptodsl/_tracing/session.py +++ b/ptodsl/ptodsl/_tracing/session.py @@ -164,6 +164,8 @@ def __init__(self, module_spec, module, entry_function): } self._subkernel_stack: list[SubkernelTraceFrame] = [] self._carry_loop_stack = [] + self._active_physical_sections: set[tuple[str, str]] = set() + self._authored_physical_sections: dict[tuple[str, str], set[str]] = {} self._inline_subkernel_counter = 0 self._escaped_inline_values: dict[object, tuple[str, str]] = {} @@ -216,6 +218,70 @@ def bind_entry_block(self, entry_block) -> None: """Record the root entry block for the active trace.""" self.entry_block = entry_block + def _physical_section_function_key(self) -> tuple[str, str]: + return (self.current_function_owner_symbol_name, self.current_function_symbol_name) + + def _enclosing_physical_section(self): + owner = InsertionPoint.current.block.owner + while owner is not None: + if owner.operation.name in {"pto.section.cube", "pto.section.vector"}: + return owner + owner = owner.operation.parent + return None + + def _existing_physical_section_kinds(self) -> set[str]: + kinds = set() + for op in self.current_function.body.blocks[0].operations: + if op.operation.name == "pto.section.cube": + kinds.add("cube") + elif op.operation.name == "pto.section.vector": + kinds.add("vector") + return kinds + + @contextmanager + def enter_physical_section(self, kind: str): + """Create one explicit cube/vector section in the active function.""" + module_spec = self.current_function_module_spec + if ( + getattr(module_spec, "kernel_kind_explicit", False) + and getattr(module_spec, "kernel_kind", None) in {"cube", "vector"} + ): + raise RuntimeError( + "pto.section() cannot be combined with explicit " + "@pto.jit(kernel_kind=...); use exactly one physical-placement contract" + ) + + function_key = self._physical_section_function_key() + if ( + function_key in self._active_physical_sections + or self._enclosing_physical_section() is not None + ): + raise RuntimeError("nested pto.section() scopes are not allowed") + + authored_kinds = self._authored_physical_sections.get(function_key, set()) + if kind in authored_kinds or kind in self._existing_physical_section_kinds(): + raise RuntimeError( + f"pto.section({kind!r}) may appear at most once in a function" + ) + + section_op = _pto.SectionCubeOp() if kind == "cube" else _pto.SectionVectorOp() + section_block = section_op.body.blocks.append() + authored_kinds = self._authored_physical_sections.setdefault(function_key, set()) + authored_kinds.add(kind) + self._active_physical_sections.add(function_key) + try: + with InsertionPoint(section_block): + yield section_op + except BaseException: + if section_op.operation.parent is not None: + section_op.operation.erase() + authored_kinds.remove(kind) + if not authored_kinds: + self._authored_physical_sections.pop(function_key, None) + raise + finally: + self._active_physical_sections.remove(function_key) + def validate_surface_value_access(self, value) -> None: """Reject inline-subkernel SSA values that escaped their outlined helper body.""" try: diff --git a/ptodsl/ptodsl/_vmi_namespace.py b/ptodsl/ptodsl/_vmi_namespace.py index b22287b23c..ccc9ada7bd 100644 --- a/ptodsl/ptodsl/_vmi_namespace.py +++ b/ptodsl/ptodsl/_vmi_namespace.py @@ -12,7 +12,7 @@ from collections.abc import Sequence from mlir.dialects import pto as _pto -from mlir.ir import BF16Type, F16Type, F32Type, Float8E4M3FNType, Float8E5M2Type, IntegerType, MemRefType +from mlir.ir import BF16Type, F16Type, F32Type, Float8E4M3FNType, Float8E5M2Type, IntegerType, MemRefType, UnitAttr from ._scalar_coercion import coerce_scalar_to_type from ._surface_values import _coerce_index_value, _try_get_constant_index, unwrap_surface_value, wrap_surface_value @@ -499,7 +499,7 @@ def _emit_reduce( ) kwargs = {"group": group, "pmode": pmode, "loc": loc, "ip": ip} if reassoc is not _UNSPECIFIED: - kwargs["reassoc"] = reassoc + kwargs["reassoc"] = UnitAttr.get() return _call_value( op_name, _derive_vmi_reduce_result_type(source, group, context=context), @@ -725,9 +725,9 @@ def vbrc(value, *, size, group=None, loc=None, ip=None): ) return _call_value("vbrc", result_type, raw_value, group=group, loc=loc, ip=ip) - vcadd = staticmethod(lambda source, mask, *, group=None, pmode=None, reassoc=_UNSPECIFIED, loc=None, ip=None: _emit_reduce("vcadd", source, mask, group=group, pmode=pmode, reassoc=reassoc, loc=loc, ip=ip)) - vcmax = staticmethod(lambda source, mask, *, group=None, pmode=None, loc=None, ip=None: _emit_reduce("vcmax", source, mask, group=group, pmode=pmode, loc=loc, ip=ip)) - vcmin = staticmethod(lambda source, mask, *, group=None, pmode=None, loc=None, ip=None: _emit_reduce("vcmin", source, mask, group=group, pmode=pmode, loc=loc, ip=ip)) + vcadd = staticmethod(lambda source, mask, *, group=1, pmode=None, reassoc=_UNSPECIFIED, loc=None, ip=None: _emit_reduce("vcadd", source, mask, group=1 if group is None else group, pmode=pmode, reassoc=reassoc, loc=loc, ip=ip)) + vcmax = staticmethod(lambda source, mask, *, group=1, pmode=None, loc=None, ip=None: _emit_reduce("vcmax", source, mask, group=1 if group is None else group, pmode=pmode, loc=loc, ip=ip)) + vcmin = staticmethod(lambda source, mask, *, group=1, pmode=None, loc=None, ip=None: _emit_reduce("vcmin", source, mask, group=1 if group is None else group, pmode=pmode, loc=loc, ip=ip)) @staticmethod def vcvt( diff --git a/ptodsl/ptodsl/pto.py b/ptodsl/ptodsl/pto.py index 9c6279d7c7..288f5d5ac6 100644 --- a/ptodsl/ptodsl/pto.py +++ b/ptodsl/ptodsl/pto.py @@ -153,7 +153,7 @@ # ── Control flow ────────────────────────────────────────────────────────────── from ._control_flow import ( # noqa: F401 - vecscope, + section, vecscope, for_, if_, yield_, static_range, LoopHandle, BranchHandle, diff --git a/ptodsl/tests/test_flash_attention_demo_compile.py b/ptodsl/tests/test_flash_attention_demo_compile.py index ef4d5ee234..063b5b812d 100644 --- a/ptodsl/tests/test_flash_attention_demo_compile.py +++ b/ptodsl/tests/test_flash_attention_demo_compile.py @@ -61,11 +61,11 @@ def main() -> None: wrapper_text = demo.emit_flash_attention_mlir(head_dim=128, causal=False, block_q=128, block_kv=128) expect_parse_roundtrip_and_verify(wrapper_text, "flash attention wrapper-emitted MLIR") expect("func.func @flash_attention_kernel" in wrapper_text, "wrapper compile should emit the flash_attention_kernel entry") - expect(wrapper_text.count("module") >= 2, "wrapper compile should emit an outer container plus child modules") + expect(wrapper_text.count("module") == 1, "wrapper compile should leave mixed VPTO input unpartitioned") expect( 'pto.backend = "vpto"' in wrapper_text and 'pto.target_arch = "a5"' in wrapper_text - and 'pto.kernel_kind = #pto.kernel_kind' in wrapper_text, + and 'pto.kernel_kind' not in wrapper_text, "flash attention wrapper compile should encode the VPTO backend directly on the child module", ) expect("func.func @materialize_tile_bounds" in wrapper_text, "wrapper compile should emit the SIMT helper function") @@ -94,11 +94,11 @@ def main() -> None: specialized_text = compiled.mlir_text() expect_parse_roundtrip_and_verify(specialized_text, "flash attention specialized MLIR") expect("func.func @flash_attention_kernel" in specialized_text, "direct compile should emit the flash_attention_kernel entry") - expect(specialized_text.count("module") >= 2, "direct compile should keep the backend-partitioned container shape") + expect(specialized_text.count("module") == 1, "direct compile should keep the unpartitioned VPTO module shape") expect( 'pto.backend = "vpto"' in specialized_text and 'pto.target_arch = "a5"' in specialized_text - and 'pto.kernel_kind = #pto.kernel_kind' in specialized_text, + and 'pto.kernel_kind' not in specialized_text, "direct compile should encode the VPTO backend directly on the child module", ) expect("!pto.tile_buf" in block64_text, "BLOCK=64 specialization MLIR missing specialized tile") expect("pto.entry" in default_text, "default @pto.jit entry child should carry the explicit entry marker") expect("pto.entry" in explicit_text, "explicit @pto.jit entry child should carry the explicit entry marker") - expect(default_text.count("module") >= 2, "default @pto.jit should emit an outer container plus one child module") - expect(block64_text.count("module") >= 2, "specialized @pto.jit should keep the outer-plus-child container shape") - expect('module attributes {pto.target_arch = "a5"}' in default_text, "outer container should carry only shared target-arch metadata") + expect(default_text.count("module") == 1, "default @pto.jit should leave an unspecified-kind VPTO module unpartitioned") + expect(block64_text.count("module") == 1, "specialized @pto.jit should keep the unpartitioned module shape") + expect( + 'module attributes {pto.backend = "vpto", pto.target_arch = "a5"}' in default_text, + "unpartitioned VPTO module should carry backend and target metadata", + ) expect('pto.mode = ' not in default_text, "generated PTODSL container IR should no longer expose public pto.mode") expect( 'pto.backend = "vpto"' in default_text and 'pto.target_arch = "a5"' in default_text - and 'pto.kernel_kind = #pto.kernel_kind' in default_text, + and 'pto.kernel_kind' not in default_text, "primary VPTO child module should carry PTOAS-facing backend metadata directly on the child module", ) expect( 'pto.backend = "vpto"' in explicit_text and 'pto.target_arch = "a5"' in explicit_text - and 'pto.kernel_kind = #pto.kernel_kind' in explicit_text, + and 'pto.kernel_kind' not in explicit_text, "explicit specialization child module should keep the same VPTO child metadata shape", ) expect( @@ -4370,8 +4379,8 @@ def inline_source_backed_probe(ptr: pto.ptr(pto.f32, "gm"), rows: pto.i32): "@pto.jit(entry=False) handles should expose an explicit, stable cache-signature protocol", ) expect( - helper_cache_signature[7] == "vector" and helper_cache_signature[8] is False, - "default @pto.jit handles should keep vector as the effective kernel kind while recording that it was not explicit", + helper_cache_signature[7] is None and helper_cache_signature[8] is False, + "default @pto.jit handles should leave the effective kernel kind unspecified", ) expect_raises( RuntimeError, @@ -4435,7 +4444,7 @@ def inline_source_backed_probe(ptr: pto.ptr(pto.f32, "gm"), rows: pto.i32): ) expect( kernel_module_call_text.count('pto.backend = "vpto"') >= 2 - and kernel_module_call_text.count('pto.kernel_kind = #pto.kernel_kind') >= 2, + and 'pto.kernel_kind' not in kernel_module_call_text, "entry-plus-helper specialization should materialize separate child modules for caller and callee", ) ast_rewrite_kernel_module_text = entry_calls_ast_rewrite_kernel_module_probe.compile().mlir_text() @@ -4480,7 +4489,7 @@ def inline_source_backed_probe(ptr: pto.ptr(pto.f32, "gm"), rows: pto.i32): expect( 'pto.backend = "vpto"' in mixed_backend_text and 'pto.target_arch = "a5"' in mixed_backend_text - and 'pto.kernel_kind = #pto.kernel_kind' in mixed_backend_text, + and 'pto.kernel_kind' not in mixed_backend_text, "mixed-backend callee child should preserve the callee's VPTO backend through child pto.backend metadata", ) expect( @@ -4653,7 +4662,10 @@ def fake_compile_launch_cpp( export_macro, ): expect(launch_cpp.is_file(), "native build should materialize launch.cpp before compiling it") - expect(kernel_kind in {"vector", "cube"}, "native build should forward the authored kernel kind") + expect( + kernel_kind in {None, "vector", "cube"}, + "native build should forward the optional authored kernel kind", + ) launch_target_arches.append(target_arch) expect(export_macro.endswith("_EXPORTS"), "native build should preserve launch export macro naming") launch_object.write_text("fake launch object\n", encoding="utf-8") @@ -4661,7 +4673,10 @@ def fake_compile_launch_cpp( def fake_link_shared_library(launch_object, kernel_object, shared_library, *, kernel_kind): expect(launch_object.is_file(), "native build should compile launch.cpp before linking") expect(kernel_object.is_file(), "native build should run ptoas before shared-library link") - expect(kernel_kind in {"vector", "cube"}, "native build should preserve kernel-kind-aware link flags") + expect( + kernel_kind in {None, "vector", "cube"}, + "native build should preserve the optional kernel kind", + ) shared_library.write_text("fake shared library\n", encoding="utf-8") with mock.patch.object(native_build_runtime, "artifact_paths", side_effect=fake_artifacts), mock.patch.object( @@ -4726,9 +4741,10 @@ def fake_link_shared_library(launch_object, kernel_object, shared_library, *, ke f"{label} native build should hand the backend-partitioned container MLIR to ptoas unchanged", ) if module_spec.jit_source is None: + expected_module_count = 1 if module_spec.kernel_kind is None else 2 expect( - observation["mlir_text"].count("module") >= 2, - f"{label} native build should route the unified outer+child container through ptoas", + observation["mlir_text"].count("module") >= expected_module_count, + f"{label} native build should route the authored module shape through ptoas", ) with TemporaryDirectory() as tmpdir: tmpdir_path = Path(tmpdir) @@ -4869,11 +4885,10 @@ def fake_run_ptoas_cmd(cmd, *, cwd=None): and 'func.func @host_vec_copy_explicit(' in merged_cross_mode_text, "merge_jit_modules() should no longer reject child modules that differ only in compile policy", ) - merged_same_mode_text = str(pto.merge_jit_modules(host_vec_copy.compile(), host_vec_copy.compile(BLOCK=64))) - expect_parse_roundtrip_and_verify(merged_same_mode_text, "merged same-mode PTODSL container") - expect( - merged_same_mode_text.count('func.func @host_vec_copy(') == 2, - "merge_jit_modules() should preserve both primary child modules in the merged container", + expect_raises( + ValueError, + lambda: pto.merge_jit_modules(host_vec_copy.compile(), host_vec_copy.compile(BLOCK=64)), + "cannot merge duplicate symbols into a flat module", ) runtime_metadata_text = runtime_metadata_kernel.compile().mlir_text() @@ -6372,6 +6387,15 @@ def _enter_inline_simt_with_resource_attr(): op_name in vmi_wrapper_dispatch_text, f"representative {op_name} wrapper dispatch should emit the matching generated VMI op", ) + for op_name, expected_count in (("vcadd", 2), ("vcmax", 3), ("vcmin", 2)): + expect( + vmi_wrapper_dispatch_text.count(f"pto.vmi.{op_name}") == expected_count, + f"pto.vmi.{op_name} should emit both omitted-group and explicit-group probes", + ) + expect( + vmi_wrapper_dispatch_text.count("group = 1") >= 6, + "VMI reductions should make the omitted group equivalent to group=1", + ) expect( vmi_wrapper_dispatch_text.count("pto.vmi.vload") == 7, "vmi wrapper dispatch probe should lower seven explicit VMI loads", diff --git a/ptodsl/tests/test_runtime_toolchain.py b/ptodsl/tests/test_runtime_toolchain.py index bea4a0fcc3..dae36bb544 100644 --- a/ptodsl/tests/test_runtime_toolchain.py +++ b/ptodsl/tests/test_runtime_toolchain.py @@ -21,6 +21,7 @@ def test_a2_a3_use_c220_aicore_arch(self): self.assertEqual( toolchain.aicore_arch_for_kernel_kind("vector", "a3"), "dav-c220-vec" ) + self.assertEqual(toolchain.aicore_arch_for_kernel_kind(None, "a3"), "dav-c220") self.assertEqual( toolchain.aicore_arch_for_kernel_kind("cube", "a3"), "dav-c220-cube" ) @@ -35,6 +36,7 @@ def test_a5_uses_c310_aicore_arch(self): self.assertEqual( toolchain.aicore_arch_for_kernel_kind("cube", "a5"), "dav-c310-cube" ) + self.assertEqual(toolchain.aicore_arch_for_kernel_kind(None, "a5"), "dav-c310") def test_native_launch_flags_use_target_arch(self): with mock.patch.object(native_build, "common_include_flags", return_value=[]): diff --git a/ptodsl/tests/test_section.py b/ptodsl/tests/test_section.py new file mode 100644 index 0000000000..aa8a34cac3 --- /dev/null +++ b/ptodsl/tests/test_section.py @@ -0,0 +1,141 @@ +#!/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. + +"""Focused tracing coverage for explicit physical section hints.""" + +from ptodsl import pto +from ptodsl._bootstrap import make_context +from mlir.ir import Module + + +@pto.jit(target="a5", mode="explicit", ast_rewrite=False) +def explicit_sections_probe(): + event_id = pto.const(1, dtype=pto.i32) + with pto.section("cube"): + pto.set_flag("MTE2", "S", event_id=0) + pto.wait_flag("S", "MTE2", event_id=event_id) + with pto.section("vector"): + pto.wait_flag("MTE2", "S", event_id=0) + pto.set_flag("S", "MTE2", event_id=event_id) + + +@pto.jit(target="a5", mode="explicit", ast_rewrite=False) +def duplicate_section_probe(): + with pto.section("cube"): + pass + with pto.section("cube"): + pass + + +@pto.jit(target="a5", mode="explicit", ast_rewrite=False) +def nested_section_probe(): + with pto.section("cube"): + with pto.section("vector"): + pass + + +@pto.jit( + target="a5", + mode="explicit", + kernel_kind="cube", + ast_rewrite=False, +) +def kernel_kind_section_probe(): + with pto.section("cube"): + pass + + +@pto.jit(target="a5", ast_rewrite=False) +def auto_mode_section_probe(): + with pto.section("vector"): + pass + + +@pto.jit( + target="a5", + mode="explicit", + kernel_kind=None, + ast_rewrite=False, +) +def unspecified_kernel_kind_section_probe(): + with pto.section("vector"): + pto.wait_flag("MTE2", "S", event_id=0) + + +@pto.jit(target="a5", mode="explicit", ast_rewrite=False) +def recovered_section_probe(): + try: + with pto.section("cube"): + raise RuntimeError("abort authored section") + except RuntimeError: + pass + with pto.section("cube"): + pto.set_flag("MTE2", "S", event_id=0) + + +def _expect_raises(exc_type, callback, message): + try: + callback() + except exc_type as exc: + assert message in str(exc), str(exc) + return + raise AssertionError(f"expected {exc_type.__name__} containing {message!r}") + + +def main() -> None: + text = explicit_sections_probe.compile().mlir_text() + assert text.count("pto.section.cube {") == 1 + assert text.count("pto.section.vector {") == 1 + assert "pto.set_flag[, , ]" in text + assert "pto.wait_flag_dyn[, " in text + assert "pto.wait_flag[, , ]" in text + assert "pto.set_flag_dyn[, " in text + with make_context() as context: + module = Module.parse(text, context) + module.operation.verify() + + recovered_text = recovered_section_probe.compile().mlir_text() + assert recovered_text.count("pto.section.cube {") == 1 + assert "pto.set_flag[, , ]" in recovered_text + + _expect_raises(TypeError, lambda: pto.section(1), "expects 'cube' or 'vector'") + _expect_raises(ValueError, lambda: pto.section("scalar"), "expects 'cube' or 'vector'") + _expect_raises(ValueError, lambda: pto.section("CUBE"), "expects 'cube' or 'vector'") + _expect_raises( + RuntimeError, + lambda: duplicate_section_probe.compile(), + "may appear at most once in a function", + ) + _expect_raises( + RuntimeError, + lambda: nested_section_probe.compile(), + "nested pto.section() scopes are not allowed", + ) + _expect_raises( + RuntimeError, + lambda: kernel_kind_section_probe.compile(), + "cannot be combined with explicit @pto.jit(kernel_kind=...)", + ) + _expect_raises( + RuntimeError, + lambda: auto_mode_section_probe.compile(), + "only available in @pto.jit(mode=\"explicit\")", + ) + _expect_raises( + RuntimeError, + lambda: pto.section("cube").__enter__(), + "may only be used while tracing", + ) + + unspecified_text = unspecified_kernel_kind_section_probe.compile().mlir_text() + assert "pto.section.vector {" in unspecified_text + + +if __name__ == "__main__": + main() diff --git a/ptodsl/tests/test_vector_cube_ops.py b/ptodsl/tests/test_vector_cube_ops.py index 71b975ab34..becba9215f 100644 --- a/ptodsl/tests/test_vector_cube_ops.py +++ b/ptodsl/tests/test_vector_cube_ops.py @@ -705,8 +705,8 @@ def test_sync_facades_reject_illegal_pipe_endpoints(self): cases = [ (_ops.set_cross_flag, (pto.Pipe.V, 0), "set_cross_flag(pipe, event_id)", "", ""), (_ops.wait_cross_flag, (pto.Pipe.MTE3, 0), "wait_cross_flag(pipe, event_id)", "", ""), - (_ops.set_intra_flag, (pto.Pipe.V, 0), "set_intra_flag(pipe, event_id)", ", ", ""), - (_ops.wait_intra_flag, (pto.Pipe.MTE2, 0), "wait_intra_flag(pipe, event_id)", ", , ", ""), + (_ops.set_intra_flag, ("M", 0), "set_intra_flag(pipe, event_id)", ", , , , ", ""), + (_ops.wait_intra_flag, (pto.Pipe.MTE3, 0), "wait_intra_flag(pipe, event_id)", ", ", ""), ] with patch.object(_ops._pto, "sync_set") as sync_set_op, \ diff --git a/test/lit/pto/sync_set_a5_emitc_physical_ids.pto b/test/lit/pto/sync_set_a5_emitc_physical_ids.pto index 6f23f96906..651d8daeb9 100644 --- a/test/lit/pto/sync_set_a5_emitc_physical_ids.pto +++ b/test/lit/pto/sync_set_a5_emitc_physical_ids.pto @@ -12,6 +12,7 @@ module attributes {pto.target_arch = "a5"} { func.func @sync_set_a5_emitc_physical_ids() attributes {pto.kernel} { pto.sync.set , 0 pto.sync.set , 16 + pto.sync.wait , 0 return } } @@ -19,4 +20,5 @@ module attributes {pto.target_arch = "a5"} { // CHECK-LABEL: AICORE void sync_set_a5_emitc_physical_ids() // CHECK: set_intra_block(PIPE_FIX, 0); // CHECK-NEXT: set_intra_block(PIPE_FIX, 16); +// CHECK: wait_intra_block(PIPE_V, 0); // CHECK-NOT: set_intra_block(PIPE_FIX, 32); diff --git a/test/lit/vmi_new/vmi_lane_stride_dense_load_store.pto b/test/lit/vmi_new/vmi_lane_stride_dense_load_store.pto index cc76b8e1d7..50638e061c 100644 --- a/test/lit/vmi_new/vmi_lane_stride_dense_load_store.pto +++ b/test/lit/vmi_new/vmi_lane_stride_dense_load_store.pto @@ -100,6 +100,24 @@ module { return %compact : !pto.vmi.vreg<64xf16, #pto.vmi.layout> } + func.func @vmi_ensure_lane_stride_to_contiguous_bf16_128( + %value: !pto.vmi.vreg<128xbf16, #pto.vmi.layout>) + -> !pto.vmi.vreg<128xbf16, #pto.vmi.layout> { + %compact = pto.vmi.ensure_layout %value + : !pto.vmi.vreg<128xbf16, #pto.vmi.layout> + -> !pto.vmi.vreg<128xbf16, #pto.vmi.layout> + return %compact : !pto.vmi.vreg<128xbf16, #pto.vmi.layout> + } + + func.func @vmi_ensure_lane_stride_to_contiguous_bf16_256( + %value: !pto.vmi.vreg<256xbf16, #pto.vmi.layout>) + -> !pto.vmi.vreg<256xbf16, #pto.vmi.layout> { + %compact = pto.vmi.ensure_layout %value + : !pto.vmi.vreg<256xbf16, #pto.vmi.layout> + -> !pto.vmi.vreg<256xbf16, #pto.vmi.layout> + return %compact : !pto.vmi.vreg<256xbf16, #pto.vmi.layout> + } + func.func @vmi_ensure_contiguous_to_lane_stride4_ui8( %value: !pto.vmi.vreg<64xui8, #pto.vmi.layout>) -> !pto.vmi.vreg<64xui8, #pto.vmi.layout> { @@ -182,6 +200,32 @@ module { // LOWER-NOT: pto.vmi. // LOWER-NOT: !pto.vmi. +// LOWER-LABEL: func.func @vmi_ensure_lane_stride_to_contiguous_bf16_128( +// LOWER-SAME: %[[LO:[^,]+]]: !pto.vreg<128xbf16>, %[[HI:[^)]+]]: !pto.vreg<128xbf16> +// LOWER: %[[LO_BITS:.*]] = pto.vbitcast %[[LO]] +// LOWER: %[[LO_PACKED:.*]] = pto.vpack %[[LO_BITS]], "LOWER" +// LOWER: %[[HI_BITS:.*]] = pto.vbitcast %[[HI]] +// LOWER: %[[HI_PACKED:.*]] = pto.vpack %[[HI_BITS]], "HIGHER" +// LOWER: %[[MERGED:.*]] = pto.vor %[[LO_PACKED]], %[[HI_PACKED]], +// LOWER: pto.vbitcast %[[MERGED]] +// LOWER-NOT: pto.vmi. +// LOWER-NOT: !pto.vmi. + +// LOWER-LABEL: func.func @vmi_ensure_lane_stride_to_contiguous_bf16_256( +// LOWER-SAME: %[[P0:[^,]+]]: !pto.vreg<128xbf16>, %[[P1:[^,]+]]: !pto.vreg<128xbf16>, %[[P2:[^,]+]]: !pto.vreg<128xbf16>, %[[P3:[^)]+]]: !pto.vreg<128xbf16> +// LOWER: %[[P0_BITS:.*]] = pto.vbitcast %[[P0]] +// LOWER: %[[P0_PACKED:.*]] = pto.vpack %[[P0_BITS]], "LOWER" +// LOWER: %[[P1_BITS:.*]] = pto.vbitcast %[[P1]] +// LOWER: %[[P1_PACKED:.*]] = pto.vpack %[[P1_BITS]], "HIGHER" +// LOWER: pto.vor %[[P0_PACKED]], %[[P1_PACKED]], +// LOWER: %[[P2_BITS:.*]] = pto.vbitcast %[[P2]] +// LOWER: %[[P2_PACKED:.*]] = pto.vpack %[[P2_BITS]], "LOWER" +// LOWER: %[[P3_BITS:.*]] = pto.vbitcast %[[P3]] +// LOWER: %[[P3_PACKED:.*]] = pto.vpack %[[P3_BITS]], "HIGHER" +// LOWER: pto.vor %[[P2_PACKED]], %[[P3_PACKED]], +// LOWER-NOT: pto.vmi. +// LOWER-NOT: !pto.vmi. + // LOWER-LABEL: func.func @vmi_ensure_contiguous_to_lane_stride4_ui8( // LOWER: pto.vzunpack {{.*}} : !pto.vreg<256xui8> -> !pto.vreg<128xui16> // LOWER: pto.vzunpack {{.*}} : !pto.vreg<128xui16> -> !pto.vreg<64xui32> diff --git a/test/lit/vpto/intra_block_sync_vpto_llvm.pto b/test/lit/vpto/intra_block_sync_vpto_llvm.pto index a2b7d001d3..d84d56d1e0 100644 --- a/test/lit/vpto/intra_block_sync_vpto_llvm.pto +++ b/test/lit/vpto/intra_block_sync_vpto_llvm.pto @@ -16,8 +16,22 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %event_id return } + + func.func @intra_block_sync_mte_pipes() attributes {pto.kernel} { + pto.sync.set , 3 + pto.sync.wait , 3 + return + } } // CHECK-LABEL: llvm.func @intra_block_sync_static_dynamic_mix_aiv // CHECK: llvm.call @llvm.hivm.SET.INTRA.BLOCK.mode // CHECK: llvm.call @llvm.hivm.WAIT.INTRA.BLOCK.mode + +// CHECK-LABEL: llvm.func @intra_block_sync_mte_pipes_mix_aiv +// CHECK: llvm.mlir.constant(3 : i64) : i64 +// CHECK: llvm.mlir.constant(3 : i64) : i64 +// CHECK: llvm.call @llvm.hivm.SET.INTRA.BLOCK.mode +// CHECK: llvm.mlir.constant(4 : i64) : i64 +// CHECK: llvm.mlir.constant(3 : i64) : i64 +// CHECK: llvm.call @llvm.hivm.WAIT.INTRA.BLOCK.mode diff --git a/test/lit/vpto/normalize_explicit_shared_pipe_sections.pto b/test/lit/vpto/normalize_explicit_shared_pipe_sections.pto new file mode 100644 index 0000000000..be6e6fd278 --- /dev/null +++ b/test/lit/vpto/normalize_explicit_shared_pipe_sections.pto @@ -0,0 +1,30 @@ +// 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. + +// Guard: explicit physical sections resolve shared-pipe synchronization that +// automatic section inference cannot classify uniquely. +// RUN: pto-test-opt --pto-normalize-uncovered-tile-sections %s | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + func.func @explicit_shared_pipe_sync() attributes {pto.kernel} { + pto.section.cube { + pto.set_flag[, , ] + } + pto.section.vector { + pto.wait_flag[, , ] + } + return + } +} + +// CHECK: pto.section.cube { +// CHECK-NEXT: pto.set_flag[, , ] +// CHECK-NEXT: } +// CHECK: pto.section.vector { +// CHECK-NEXT: pto.wait_flag[, , ] +// CHECK-NEXT: } diff --git a/test/lit/vpto/normalize_uncovered_micro_sections_mixed.pto b/test/lit/vpto/normalize_uncovered_micro_sections_mixed.pto new file mode 100644 index 0000000000..89e617c068 --- /dev/null +++ b/test/lit/vpto/normalize_uncovered_micro_sections_mixed.pto @@ -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. + +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto %s -o - | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.backend = "vpto", pto.target_arch = "a5"} { + func.func @mixed_micro_kernel(%src: !pto.ptr) attributes {pto.entry} { + %c0_i64 = arith.constant 0 : i64 + %l1 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %c16_i64 = arith.constant 16 : i64 + %c1_i64 = arith.constant 1 : i64 + %false = arith.constant false + pto.set_flag[, , ] + pto.sync.set , 5 + pto.mte_gm_l1_frac %src, %l1, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c16_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), + ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, + ctrl i64, i1 + pto.sync.wait , 0 + pto.wait_flag[, , ] + pto.sync.set , 0 + pto.set_flag[, , ] + pto.set_flag[, , ] + pto.vecscope { + %c64 = arith.constant 64 : index + %mask = pto.vmi.create_mask %c64 : index -> !pto.vmi.mask<64xpred> + } + pto.sync.wait , 0 + pto.wait_flag[, , ] + pto.wait_flag[, , ] + return + } + } +} + +// CHECK: module attributes +// CHECK: module attributes +// CHECK-SAME: pto.kernel_kind = #pto.kernel_kind +// CHECK: func.func @mixed_micro_kernel +// CHECK-NOT: pto.set_flag[, +// CHECK-NOT: pto.sync.set +// CHECK-NOT: pto.sync.wait +// CHECK-NOT: pto.wait_flag[, +// CHECK: pto.sync.set , 0 +// CHECK: pto.set_flag[, , ] +// CHECK: pto.set_flag[, , ] +// CHECK: pto.vecscope { +// CHECK: pto.sync.wait , 0 +// CHECK: pto.wait_flag[, , ] +// CHECK: pto.wait_flag[, , ] +// CHECK: module attributes +// CHECK-SAME: pto.kernel_kind = #pto.kernel_kind +// CHECK: func.func @mixed_micro_kernel +// CHECK-NOT: pto.sync.set +// CHECK-NOT: pto.set_flag[, +// CHECK-NOT: pto.set_flag[, +// CHECK: pto.copy_gm_to_cbuf_multi_nd2nz +// CHECK: pto.sync.wait , 0 +// CHECK: pto.wait_flag[, , ] diff --git a/test/lit/vpto/normalize_uncovered_tile_sections_reject_residual_nested_mix.pto b/test/lit/vpto/normalize_uncovered_tile_sections_reject_residual_nested_mix.pto index 1ae41f30c7..fda54770e8 100644 --- a/test/lit/vpto/normalize_uncovered_tile_sections_reject_residual_nested_mix.pto +++ b/test/lit/vpto/normalize_uncovered_tile_sections_reject_residual_nested_mix.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // Guard: normalize must reject a top-level segment that still mixes a nested -// explicit section with sibling uncovered TileOps after inferred wrapping. +// explicit section with sibling uncovered ops after inferred wrapping. // RUN: not ptoas --pto-arch=a5 --pto-backend=vpto --emit-pto-ir %s -o - 2>&1 | FileCheck %s module attributes {pto.target_arch = "a5"} { @@ -34,6 +34,6 @@ module attributes {pto.target_arch = "a5"} { } } -// CHECK: still contains an uncovered top-level TileOp segment after section normalization -// CHECK: mixes nested explicit pto.section.* with sibling TileOps outside those sections +// CHECK: still contains an uncovered top-level op segment after section normalization +// CHECK: mixes nested explicit pto.section.* with sibling ops outside those sections // CHECK: starts at 'scf.for' diff --git a/test/lit/vpto/normalize_uncovered_tile_sections_suggest_explicit.pto b/test/lit/vpto/normalize_uncovered_tile_sections_suggest_explicit.pto new file mode 100644 index 0000000000..43c5bb01ae --- /dev/null +++ b/test/lit/vpto/normalize_uncovered_tile_sections_suggest_explicit.pto @@ -0,0 +1,22 @@ +// 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. + +// Guard: ambiguous shared-pipe synchronization should recommend the explicit +// PTODSL section escape hatch without weakening automatic inference. +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto --emit-pto-ir %s -o - 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + func.func @ambiguous_shared_pipe_sync() attributes {pto.kernel} { + pto.set_flag[, , ] + return + } +} + +// CHECK: contains an uncovered top-level op segment whose section kind cannot be inferred uniquely +// CHECK-SAME: ambiguous op(s): 'pto.set_flag' +// CHECK-SAME: wrap the ambiguous region in pto.section.cube or pto.section.vector diff --git a/tools/ptoas/CMakeLists.txt b/tools/ptoas/CMakeLists.txt index 1e2918bd5a..1adff51142 100644 --- a/tools/ptoas/CMakeLists.txt +++ b/tools/ptoas/CMakeLists.txt @@ -116,6 +116,20 @@ add_custom_target(ptoas_runtime_staging ALL VERBATIM ) +# Editable Python installs resolve the wheel bootstrap from build/python. Keep +# a wheel-shaped runtime overlay there so the same launcher contract works for +# both regular and editable installations. +if(TARGET PTODSLPackage) + add_custom_command(TARGET ptoas_runtime_staging POST_BUILD + COMMAND ${CMAKE_COMMAND} -E remove_directory + "${CMAKE_BINARY_DIR}/python/ptoas/_runtime" + COMMAND ${CMAKE_COMMAND} -E copy_directory + "${PTOAS_BUILD_RUNTIME_STAGING_DIR}" + "${CMAKE_BINARY_DIR}/python/ptoas/_runtime" + VERBATIM + ) +endif() + if(TARGET PTODSLPackage) # The build-tree launcher imports ptoas._runtime_entry from build/python. # Keep that package staging coupled to every runtime staging entrypoint,