公司动态

CANN PTO TAXPY缩放累加指令

📅 2026/7/16 14:30:25
CANN PTO TAXPY缩放累加指令
TAXPY【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isaIntroductionPerform an in-place scaled accumulation (AXPY, $a \cdot x y$) on a tile: scalesrc0byscalarand accumulate it intodst.$$ \mathrm{dst}{i,j} \leftarrow \mathrm{scalar} \cdot \mathrm{src0}{i,j} \mathrm{dst}_{i,j} $$dstis both the accumulation input ($y$) and the output, and must be initialized before the call;src0($x$) is read-only;scalar($a$) is a scalar.Math InterpretationFor each element(i, j)in the valid region:$$ \mathrm{dst}{i,j}^{\text{new}} \mathrm{scalar} \cdot \mathrm{src0}{i,j} \mathrm{dst}_{i,j}^{\text{old}} $$dst: read-modify-write (RMW). Its old value is read as the accumulation base $y$, and $\mathrm{scalar} \cdot x y$ is written back.src0: read-only, contributes element-wise ($x$).scalar: the scalar scale factor ($a$), of typeTileDataSrc::DType.Unless otherwise specified, semantics are defined over the valid region and target-dependent behavior is marked as implementation-defined.C IntrinsicsDeclared ininclude/pto/common/pto_instr.hpp.template typename TileDataDst, typename TileDataSrc, typename... WaitEvents PTO_INST RecordEvent TAXPY(TileDataDst dst, TileDataSrc src0, typename TileDataSrc::DType scalar, WaitEvents ...events);ParameterDirectionMeaningdstinput/outputAccumulation base and result tile ($y$), read-modify-write,Vecsrc0inputScaled source tile ($x$), read-only,Vec, same valid shape asdstscalarinputScalar scale factor ($a$), of typeTileDataSrc::DTypeevents...inputWait events (WaitEvents); an implicitTSYNCprecedes the opTile Sizes Data TypesFor a valid tile shape of $M \times N$:TiledtypeValid shapeTileTypeNotesdsthalforfloat$M \times N$Vec(UB)Accumulation base result (RMW)src0halforfloat$M \times N$Vec(UB)Scaled source, element-wisedstandsrc0must have identical valid row and column counts.Supported Input Dtypesdstdtypesrc0dtypescalardtypeNoteshalfhalfhalfSame-type path, directvaxpyfloatfloatfloatSame-type path, directvaxpyfloathalfhalfDiff path:src0widened to FP32 before accumulationdstandsrc0must share a dtype, ordstisfloatwhilesrc0ishalf(a half→float widening accumulation is allowed). Ahalfdstwith afloatsrc0is illegal (rejected by an in-implementationstatic_assert).Implementation NotesTAXPY runs on the vector pipeline (PIPE_V) using thevaxpy($a \cdot x y$) vector intrinsic:Same type (dstandsrc0share a dtype): loadsrc0anddstper repeat, runvaxpy(dst, src0, scalar), and store back todst; tail columns shorter than a full repeat are masked by a predicate.Diff type (dstfloat,src0half): the half data ofsrc0is widened to FP32 before accumulating (on A5 viaUNPK_B16unpack followed byvcvt; on A2/A3 handled natively byvaxpywith 4-block src / 8-block dst).On A2/A3, count mode vs. norm mode is selected based on whether the repeat-stride overflows and the relation between column count and row count, so any valid shape is covered.ConstraintsConstraintApplies toReasondstandsrc0must beTileType::Vecall targetsexecutes on UB (vector pipeline)dstandsrc0share the valid shape ($M \times N$)all targetsone-to-one element mappingdstdtype ∈ {half,float}all targetsfloating-point widths supported byvaxpydst/src0dtype equal, or (float,half)all targetsonly half→float widening accumulation alloweddstmust be initialized before the callall targetsdstis read in as the accumulation base $y$Examples// dst must be initialized first (as the accumulation base y); result: dst scalar * src0 dst TAXPY(dstTile, srcTile, scalar);Seetests/npu/a5/src/st/testcase/taxpy/(A5),tests/npu/a2a3/src/st/testcase/taxpy/(A2/A3),tests/npu/kirin9030/src/st/testcase/taxpy/(Kirin9030), andtests/cpu/st/testcase/taxpy/(CPU reference) for complete ST examples.【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isa创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考