diff --git a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp index bfe0221c24..b62c087d33 100644 --- a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp @@ -35,6 +35,7 @@ #include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h" #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h" #include "mlir/Target/LLVMIR/Export.h" +#include "llvm/ADT/DenseMap.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/StringSet.h" #include "llvm/Bitcode/BitcodeWriter.h" @@ -10793,14 +10794,223 @@ getUniqueDeviceModuleByKernelKind(ModuleOp module, FunctionKernelKind kind, return matched; } +static LogicalResult mergeDeviceModulesByKernelKind(ModuleOp module) { + ModuleOp vectorModule; + ModuleOp cubeModule; + SmallVector modulesToErase; + + for (ModuleOp child : module.getOps()) { + auto kernelKind = getKernelKind(child); + if (!kernelKind) + continue; + + ModuleOp *target = nullptr; + if (*kernelKind == FunctionKernelKind::Vector) + target = &vectorModule; + else if (*kernelKind == FunctionKernelKind::Cube) + target = &cubeModule; + else + continue; + + if (!*target) { + *target = child; + continue; + } + + llvm::StringMap targetFunctions; + for (func::FuncOp fn : (*target).getOps()) + targetFunctions[fn.getSymName()] = fn; + + SmallVector functionsToErase; + for (func::FuncOp incoming : child.getOps()) { + auto it = targetFunctions.find(incoming.getSymName()); + if (it == targetFunctions.end()) + continue; + func::FuncOp existing = it->second; + if (existing.getFunctionType() != incoming.getFunctionType()) + return incoming.emitOpError() + << "cannot merge device modules: symbol '" + << incoming.getSymName() << "' has conflicting function types"; + if (!existing.isExternal() && !incoming.isExternal()) + return incoming.emitOpError() + << "cannot merge device modules: duplicate definition of symbol '" + << incoming.getSymName() << "'"; + if (existing.isExternal() && !incoming.isExternal()) { + existing.erase(); + targetFunctions[incoming.getSymName()] = incoming; + continue; + } + functionsToErase.push_back(incoming); + } + for (func::FuncOp fn : functionsToErase) + fn.erase(); + + Block *srcBody = child.getBody(); + Block *dstBody = (*target).getBody(); + while (!srcBody->empty()) { + Operation &op = srcBody->front(); + op.moveBefore(dstBody, dstBody->end()); + } + modulesToErase.push_back(child); + } + + for (ModuleOp child : modulesToErase) + child.erase(); + return success(); +} + +static void removeDuplicateDeclarations(ModuleOp module) { + DenseMap definitions; + for (func::FuncOp funcOp : module.getOps()) + if (!funcOp.isExternal()) + definitions[funcOp.getSymName()] = funcOp; + + SmallVector toErase; + for (func::FuncOp funcOp : module.getOps()) + if (funcOp.isExternal() && definitions.count(funcOp.getSymName())) + toErase.push_back(funcOp); + + for (func::FuncOp funcOp : toErase) + funcOp.erase(); +} + +static LogicalResult mergeAndCleanDuplicates(ModuleOp module) { + if (failed(mergeDeviceModulesByKernelKind(module))) + return failure(); + for (ModuleOp child : module.getOps()) + removeDuplicateDeclarations(child); + + // After merging, the vector module's entry function may contain func.call + // to a cube sub-kernel that now lives in the separate cube module. The AIV + // cannot execute cube code, so erase those calls and promote the cube + // callee to an entry function that the CANN host stub launches concurrently + // on the AIC. + ModuleOp cubeModule; + ModuleOp vectorModule; + for (ModuleOp child : module.getOps()) { + auto kind = getKernelKind(child); + if (!kind) + continue; + if (*kind == FunctionKernelKind::Cube) + cubeModule = child; + else if (*kind == FunctionKernelKind::Vector) + vectorModule = child; + } + if (!cubeModule || !vectorModule) + return success(); + + // Collect cube function symbol names that have definitions. + DenseSet cubeDefinedSymbols; + for (func::FuncOp fn : cubeModule.getOps()) + if (!fn.isExternal()) + cubeDefinedSymbols.insert(fn.getSymName()); + + // In the vector module, find func.call ops to cube-defined symbols. + // Erase those calls and mark the cube callee as an entry function. + DenseSet cubeCallsToPromote; + vectorModule.walk([&](func::CallOp call) { + StringRef callee = call.getCallee(); + if (cubeDefinedSymbols.count(callee)) + cubeCallsToPromote.insert(callee); + }); + + if (cubeCallsToPromote.empty()) + return success(); + + func::FuncOp vectorEntry; + for (func::FuncOp fn : vectorModule.getOps()) { + if (!pto::isPTOEntryFunction(fn)) + continue; + if (vectorEntry) + return fn.emitOpError() + << "cross-kernel call promotion requires one vector entry per " + "merged device module"; + vectorEntry = fn; + } + if (!vectorEntry) + return vectorModule.emitError() + << "cross-kernel call promotion requires a vector entry function"; + if (cubeCallsToPromote.size() != 1) + return vectorEntry.emitOpError() + << "cross-kernel call promotion supports exactly one cube callee"; + func::FuncOp cubeCallee; + for (func::FuncOp fn : cubeModule.getOps()) + if (!fn.isExternal() && cubeCallsToPromote.count(fn.getSymName())) + cubeCallee = fn; + if (!cubeCallee || + cubeCallee.getFunctionType() != vectorEntry.getFunctionType()) + return vectorEntry.emitOpError() + << "cross-kernel cube callee must have the same signature as its " + "vector entry"; + + // Erase func.call ops in the vector module that target cube callees. + SmallVector callsToErase; + vectorModule.walk([&](func::CallOp call) { + if (cubeCallsToPromote.count(call.getCallee())) + callsToErase.push_back(call); + }); + if (callsToErase.size() != 1) + return vectorEntry.emitOpError() + << "cross-kernel call promotion requires exactly one call to the " + "cube callee"; + for (func::CallOp call : callsToErase) { + if (call->getParentOp() != vectorEntry || call.getNumResults() != 0) + return call.emitOpError() + << "cross-kernel cube call must be unconditional and return no " + "values"; + if (call.getNumOperands() != vectorEntry.getNumArguments()) + return call.emitOpError() + << "cross-kernel cube call must forward every entry argument"; + for (auto [operand, argument] : + llvm::zip_equal(call.getOperands(), vectorEntry.getArguments())) + if (operand != argument) + return call.emitOpError() + << "cross-kernel cube call operands must be the entry arguments " + "in order"; + } + for (func::CallOp call : callsToErase) + call.erase(); + + // Erase the now-unused external declarations of cube callees in the vector module. + SmallVector declsToErase; + for (func::FuncOp fn : vectorModule.getOps()) + if (fn.isExternal() && cubeCallsToPromote.count(fn.getSymName())) + declsToErase.push_back(fn); + for (func::FuncOp fn : declsToErase) + fn.erase(); + + // Promote cube callees to entry functions and give them the same logical + // name as the vector entry so the host stub creates a single __global__ + // kernel that launches both AIC and AIV concurrently. + // Find the logical name of the vector entry function. + StringRef vectorEntryLogicalName; + for (func::FuncOp fn : vectorModule.getOps()) { + if (pto::isPTOEntryFunction(fn)) { + vectorEntryLogicalName = pto::getPTODSLLogicalNameOrSymbolName(fn); + break; + } + } + if (vectorEntryLogicalName.empty()) + return vectorEntry.emitOpError() + << "cross-kernel vector entry requires a non-empty logical name"; + + for (func::FuncOp fn : cubeModule.getOps()) { + if (!cubeCallsToPromote.count(fn.getSymName())) + continue; + // Mark as entry. + fn->setAttr(pto::kPTOEntryAttrName, UnitAttr::get(fn.getContext())); + // Set the logical name to match the vector entry. + fn->setAttr(pto::kPTODSLLogicalNameAttrName, + StringAttr::get(fn.getContext(), vectorEntryLogicalName)); + } + return success(); +} + static LogicalResult renameKernelFunctionsForKernelKind(ModuleOp module, llvm::raw_ostream &diagOS) { auto kernelKind = getKernelKind(module); - if (!kernelKind) { - diagOS << "VPTO LLVM emission failed: device module missing " - << FunctionKernelKindAttr::name << "\n"; - return failure(); - } + if (!kernelKind) + return success(); StringRef suffix; if (*kernelKind == FunctionKernelKind::Vector) @@ -10816,9 +11026,16 @@ static LogicalResult renameKernelFunctionsForKernelKind(ModuleOp module, for (func::FuncOp funcOp : module.getOps()) { if (!pto::hasExplicitPTOEntryAttr(funcOp)) continue; - if (funcOp.getSymName().ends_with(suffix)) + llvm::StringRef baseNameRef = pto::getPTODSLLogicalNameOrSymbolName(funcOp); + std::string baseName = baseNameRef.str(); + for (auto s : {kVectorSuffix, kCubeSuffix}) { + if (baseNameRef.ends_with(s)) + baseName = baseName.substr(0, baseName.size() - s.size()); + } + std::string newName = baseName + suffix.str(); + if (funcOp.getSymName() == newName) continue; - funcOp.setSymName((funcOp.getSymName() + suffix).str()); + funcOp.setSymName(newName); } return success(); } @@ -10970,6 +11187,11 @@ static LogicalResult runPipeline(ModuleOp module, llvm::raw_ostream &diagOS, OwningOpRef clonedOp(module->clone()); ModuleOp clonedModule = cast(*clonedOp); + if (failed(mergeAndCleanDuplicates(clonedModule))) { + diagOS << "VPTO LLVM emission failed: unable to merge device modules\n"; + return failure(); + } + if (failed(validateVPTOAuthoringIR(clonedModule, &diagOS))) { diagOS << "VPTO LLVM emission failed: authoring-stage VPTO legality " "validation failed\n"; diff --git a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp index 3ba307ffa2..a62b4789af 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp @@ -11631,11 +11631,8 @@ static void mergeDeviceModulesByKernelKind(ModuleOp module) { static LogicalResult renameKernelFunctionsForKernelKind(ModuleOp module, llvm::raw_ostream &diagOS) { auto kernelKind = getKernelKind(module); - if (!kernelKind) { - diagOS << "VPTO LLVM emission failed: device module missing " - << FunctionKernelKindAttr::name << "\n"; - return failure(); - } + if (!kernelKind) + return success(); StringRef suffix; if (*kernelKind == FunctionKernelKind::Vector) @@ -11651,9 +11648,16 @@ static LogicalResult renameKernelFunctionsForKernelKind(ModuleOp module, for (func::FuncOp funcOp : module.getOps()) { if (!pto::hasExplicitPTOEntryAttr(funcOp)) continue; - if (funcOp.getSymName().ends_with(suffix)) + llvm::StringRef baseNameRef = pto::getPTODSLLogicalNameOrSymbolName(funcOp); + std::string baseName = baseNameRef.str(); + for (auto s : {kVectorSuffix, kCubeSuffix}) { + if (baseNameRef.ends_with(s)) + baseName = baseName.substr(0, baseName.size() - s.size()); + } + std::string newName = baseName + suffix.str(); + if (funcOp.getSymName() == newName) continue; - funcOp.setSymName((funcOp.getSymName() + suffix).str()); + funcOp.setSymName(newName); } return success(); } diff --git a/lib/PTO/Transforms/VPTONormalizeContainer.cpp b/lib/PTO/Transforms/VPTONormalizeContainer.cpp index d5b22af3d6..2d6f7edf5a 100644 --- a/lib/PTO/Transforms/VPTONormalizeContainer.cpp +++ b/lib/PTO/Transforms/VPTONormalizeContainer.cpp @@ -38,18 +38,17 @@ static LogicalResult verifyNormalizedVPTOContainer(ModuleOp module) { "submodules"; } hasChildModules = true; - if (!isVPTOKernelSubmodule(child)) { + if (!isVPTOKernelSubmodule(child)) return child.emitError() << "expected VPTO kernel submodule to carry 'pto.kernel_kind'"; - } } if (hasChildModules) return success(); return module.emitError() - << "expected VPTO input to be a kernel submodule with " - "'pto.kernel_kind' or a container of kernel submodules"; + << "expected VPTO kernel submodule to carry 'pto.kernel_kind' " + "or a container of kernel submodules"; } struct VPTONormalizeContainerPass diff --git a/lib/PTO/Transforms/VPTOSplitCVModule.cpp b/lib/PTO/Transforms/VPTOSplitCVModule.cpp index fcea35acd6..28a33cf6c2 100644 --- a/lib/PTO/Transforms/VPTOSplitCVModule.cpp +++ b/lib/PTO/Transforms/VPTOSplitCVModule.cpp @@ -260,11 +260,24 @@ static LogicalResult splitCVModule(ModuleOp module) { if (hasKernelKind(module)) return materializeExplicitKernelKindSections(module); if (hasKernelKindChildModule(module)) { + SmallVector entryChildrenToFlatten; for (ModuleOp child : module.getOps()) { - if (!hasKernelKind(child)) - continue; - if (failed(materializeExplicitKernelKindSections(child))) - return failure(); + if (hasKernelKind(child)) { + if (failed(materializeExplicitKernelKindSections(child))) + return failure(); + } else if (hasCVSections(child)) { + if (failed(splitCVModule(child))) + return failure(); + entryChildrenToFlatten.push_back(child); + } + } + for (ModuleOp entryChild : entryChildrenToFlatten) { + SmallVector splitChildren; + for (ModuleOp gc : entryChild.getOps()) + splitChildren.push_back(gc); + for (ModuleOp gc : splitChildren) + gc->moveBefore(module.getBody(), module.getBody()->end()); + entryChild.erase(); } return success(); } diff --git a/test/lit/vpto/backend_mixed_child_missing_kernel_kind_invalid.pto b/test/lit/vpto/backend_mixed_child_missing_kernel_kind_invalid.pto new file mode 100644 index 0000000000..0d832f73bb --- /dev/null +++ b/test/lit/vpto/backend_mixed_child_missing_kernel_kind_invalid.pto @@ -0,0 +1,20 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto-llvm-ir %s -o - 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @valid_kernel() { return } + } + module { + func.func @silently_dropped_before_fix() { return } + } +} + +// CHECK: expected VPTO kernel submodule to carry 'pto.kernel_kind' diff --git a/test/lit/vpto/cann900_cross_kernel_call_invalid.pto b/test/lit/vpto/cann900_cross_kernel_call_invalid.pto new file mode 100644 index 0000000000..2790455518 --- /dev/null +++ b/test/lit/vpto/cann900_cross_kernel_call_invalid.pto @@ -0,0 +1,27 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto --cann-output-version=9.0.0 --emit-vpto-llvm-ir %s -o - 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func private @cube_peer(!pto.ptr, !pto.ptr) attributes {no_inline} + func.func @mixed_kernel(%arg0: !pto.ptr, %arg1: !pto.ptr) + attributes {pto.entry} { + call @cube_peer(%arg1, %arg0) : (!pto.ptr, !pto.ptr) -> () + return + } + } + module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @cube_peer(%arg0: !pto.ptr, %arg1: !pto.ptr) { + return + } + } +} + +// CHECK: cross-kernel cube call operands must be the entry arguments in order diff --git a/test/lit/vpto/cann900_duplicate_device_definition_invalid.pto b/test/lit/vpto/cann900_duplicate_device_definition_invalid.pto new file mode 100644 index 0000000000..5321750f8f --- /dev/null +++ b/test/lit/vpto/cann900_duplicate_device_definition_invalid.pto @@ -0,0 +1,20 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto --cann-output-version=9.0.0 --emit-vpto-llvm-ir %s -o - 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @duplicate() { return } + } + module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @duplicate() { return } + } +} + +// CHECK: {{cannot merge device modules: duplicate definition of symbol 'duplicate'|redefinition of symbol named 'duplicate'}} diff --git a/tools/ptoas/VPTOHostStubEmission.cpp b/tools/ptoas/VPTOHostStubEmission.cpp index 103e4f80eb..fbb5dc7a25 100644 --- a/tools/ptoas/VPTOHostStubEmission.cpp +++ b/tools/ptoas/VPTOHostStubEmission.cpp @@ -80,7 +80,11 @@ static LogicalResult collectVPTOKernelStubDecls( if (!pto::isPTOEntryFunction(func)) return; - std::string logicalName = getLogicalKernelName(func.getSymName()); + std::string logicalName; + if (func->hasAttr(pto::kPTODSLLogicalNameAttrName)) + logicalName = pto::getPTODSLLogicalNameOrSymbolName(func).str(); + else + logicalName = getLogicalKernelName(func.getSymName()); SmallVector argTypes; argTypes.reserve(func.getNumArguments()); for (Type type : func.getArgumentTypes())