diff --git a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp index b64e459783..2c60df94b9 100644 --- a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp @@ -2410,8 +2410,16 @@ static FailureOr buildPstuCallee(MLIRContext *context, pto::PstuOp op return failure(); } -static StringRef buildVstusCallee(MLIRContext *context) { - return StringAttr::get(context, "llvm.hivm.vstus").getValue(); +static FailureOr buildVstusCallee(MLIRContext *context, + Type valueType) { + std::string vec = + getMemoryElementTypeFragment(getElementTypeFromVectorLike(valueType)); + auto lanes = getElementCountFromVectorLike(valueType); + if (vec.empty() || !lanes) + return failure(); + return StringAttr::get(context, "llvm.hivm.vstus.v" + + std::to_string(*lanes) + vec) + .getValue(); } static FailureOr buildVstusPostCallee(MLIRContext *context, @@ -7512,7 +7520,8 @@ class LowerVstusOpPattern final : public OpConversionPattern { "unexpected converted vstus operand/result types"); } - FailureOr calleeName = buildVstusCallee(op.getContext()); + FailureOr calleeName = + buildVstusCallee(op.getContext(), op.getValue().getType()); if (usePostIntrinsic) calleeName = buildVstusPostCallee(op.getContext(), op.getValue().getType()); diff --git a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp index 32250a089f..81a3ac9bd7 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp @@ -2438,8 +2438,16 @@ static FailureOr buildPstuCallee(MLIRContext *context, pto::PstuOp op return failure(); } -static StringRef buildVstusCallee(MLIRContext *context) { - return StringAttr::get(context, "llvm.hivm.vstus").getValue(); +static FailureOr buildVstusCallee(MLIRContext *context, + Type valueType) { + std::string vec = + getMemoryElementTypeFragment(getElementTypeFromVectorLike(valueType)); + auto lanes = getElementCountFromVectorLike(valueType); + if (vec.empty() || !lanes) + return failure(); + return StringAttr::get(context, "llvm.hivm.vstus.v" + + std::to_string(*lanes) + vec) + .getValue(); } static FailureOr buildVstusPostCallee(MLIRContext *context, @@ -8111,7 +8119,8 @@ class LowerVstusOpPattern final : public OpConversionPattern { "unexpected converted vstus operand/result types"); } - FailureOr calleeName = buildVstusCallee(op.getContext()); + FailureOr calleeName = + buildVstusCallee(op.getContext(), op.getValue().getType()); if (usePostIntrinsic) calleeName = buildVstusPostCallee(op.getContext(), op.getValue().getType()); diff --git a/test/lit/vpto/vreg_low_precision_memory_vpto_llvm.pto b/test/lit/vpto/vreg_low_precision_memory_vpto_llvm.pto index 37ba0cbb21..d89f1c2cbe 100644 --- a/test/lit/vpto/vreg_low_precision_memory_vpto_llvm.pto +++ b/test/lit/vpto/vreg_low_precision_memory_vpto_llvm.pto @@ -95,7 +95,8 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind @llvm.hivm.vldas // CANN: call { <256 x i8>, <32 x i8>, ptr addrspace(6) } @llvm.hivm.vldus.v256s8 // BETA: call { <256 x i8>, <32 x i8>, ptr addrspace(6) } @llvm.hivm.vldus.v256hif8 -// CHECK: call <32 x i8> @llvm.hivm.vstus +// CANN: call <32 x i8> @llvm.hivm.vstus.v256s8 +// BETA: call <32 x i8> @llvm.hivm.vstus.v256hif8 // CHECK: call <32 x i8> @llvm.hivm.vstur // CHECK: call void @llvm.hivm.vstar diff --git a/test/lit/vpto/vstus_mixed_types_vpto_llvm.pto b/test/lit/vpto/vstus_mixed_types_vpto_llvm.pto new file mode 100644 index 0000000000..e74b5ba425 --- /dev/null +++ b/test/lit/vpto/vstus_mixed_types_vpto_llvm.pto @@ -0,0 +1,37 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software; you can redistribute it and/or modify it under the terms of +// the CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for the specific language governing permissions and limitations under the License. + +// Keep vstus callees type-specific so different element types can coexist in +// one kernel without conflicting LLVM declarations. +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto-llvm-ir %s -o - 2>&1 | FileCheck %s --check-prefix=DEFAULT +// RUN: ptoas --cann-output-version=9.0.0 --pto-arch=a5 --pto-backend=vpto --emit-vpto-llvm-ir %s -o - 2>&1 | FileCheck %s --check-prefix=CANN900 + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @vstus_mixed_types(%f32_dst: !pto.ptr, + %u8_dst: !pto.ptr, + %f32_value: !pto.vreg<64xf32>, + %u8_value: !pto.vreg<256xui8>) attributes {pto.kernel} { + %c0_i32 = arith.constant 0 : i32 + pto.vecscope { + %f32_align = pto.init_align : !pto.align + %f32_align_out = pto.vstus %f32_align, %c0_i32, %f32_value, %f32_dst : + !pto.align, i32, !pto.vreg<64xf32>, !pto.ptr -> !pto.align + pto.vstas %f32_align_out, %f32_dst, %c0_i32 : !pto.align, !pto.ptr, i32 + %u8_align = pto.init_align : !pto.align + %u8_align_out = pto.vstus %u8_align, %c0_i32, %u8_value, %u8_dst : + !pto.align, i32, !pto.vreg<256xui8>, !pto.ptr -> !pto.align + pto.vstas %u8_align_out, %u8_dst, %c0_i32 : !pto.align, !pto.ptr, i32 + } + return + } +} + +// DEFAULT-LABEL: define void @vstus_mixed_types_mix_aiv +// DEFAULT: call <32 x i8> @llvm.hivm.vstus.v64f32 +// DEFAULT: call <32 x i8> @llvm.hivm.vstus.v256i8 + +// CANN900-LABEL: define void @vstus_mixed_types_mix_aiv +// CANN900: call <32 x i8> @llvm.hivm.vstus.v64f32 +// CANN900: call <32 x i8> @llvm.hivm.vstus.v256i8