| @@ -514,63 +514,115 @@ PTO_INTERNAL void AbsReduceMax_b16_ND_2D(__ubuf__ T *srcPtr, __ubuf__ T *maxPtr, | |||
| 514 | vstas(ureg_max, writePtr, 0, POST_UPDATE); | 514 | vstas(ureg_max, writePtr, 0, POST_UPDATE); |
| 515 | (void)validCols; // padded source makes validCols implicit; retained for API symmetry | 515 | (void)validCols; // padded source makes validCols implicit; retained for API symmetry |
| 516 | } | 516 | } |
| 517 | +// Constants and registers shared by all fp32 OCP exponent/scaling extraction variants. | ||
| 518 | +// Constants and registers shared by all fp32 OCP exponent/scaling extraction variants. | ||
| 519 | +// Call InitF32OcpQuantCtx() inside a __VEC_SCOPE__ before use. | ||
| 520 | +struct F32OcpQuantCtx { | ||
| 521 | + vector_s32 vb32_b8_nan, vb32_f32_nan, vb32_b8_emax, vb32_exp_mask, vb32_mantissa_mask, vb32_exp_max; | ||
| 522 | + vector_s32 vb32_recip_min_scale, vb32_zero; | ||
| 523 | + vector_bool preg_special, preg_nan, preg_min_scale; | ||
| 524 | + static constexpr int shr = 23; | ||
| 525 | +}; | ||
| 517 | 526 | ||
| 527 | +PTO_INTERNAL void InitF32OcpQuantCtx(F32OcpQuantCtx &ctx) | ||
| 528 | +{ | ||
| 529 | + vbr(ctx.vb32_exp_mask, 0x7F800000); | ||
| 530 | + vbr(ctx.vb32_mantissa_mask, 0x007FFFFF); | ||
| 531 | + vbr(ctx.vb32_b8_nan, 0xFF); | ||
| 532 | + vbr(ctx.vb32_f32_nan, 0x7FC00000); | ||
| 533 | + vbr(ctx.vb32_exp_max, 0xFE); | ||
| 534 | + vbr(ctx.vb32_b8_emax, 8); | ||
| 535 | + vbr(ctx.vb32_recip_min_scale, 0x7F000000); | ||
| 536 | + vbr(ctx.vb32_zero, 0); | ||
| 537 | +} | ||
| 538 | + | ||
| 539 | +// Compute exponent and scaling from one VL of max data, using a pre-initialised ctx. | ||
| 540 | +// Outputs: shared_exp written to expPtr (PK4_B32), scaling written to scalingPtr (NORM). | ||
| 541 | +PTO_INTERNAL void ComputeF32OcpExpAndScaling(F32OcpQuantCtx &ctx, vector_s32 &vb32_shared_exp, vector_s32 &vb32_scaling, | ||
| 542 | + vector_f32 &vb32_max, MaskReg &preg_b32, __ubuf__ int32_t *maxPtrRaw, | ||
| 543 | + uint32_t loadOff) | ||
| 544 | +{ | ||
| 545 | + vlds((vector_s32 &)vb32_max, maxPtrRaw, loadOff, NORM); | ||
| 546 | + vector_s32 vb32_exponent, vb32_mantissa; | ||
| 547 | + vand((vector_s32 &)vb32_exponent, (vector_s32 &)vb32_max, ctx.vb32_exp_mask, preg_b32, MODE_ZEROING); | ||
| 548 | + vand((vector_s32 &)vb32_mantissa, (vector_s32 &)vb32_max, ctx.vb32_mantissa_mask, preg_b32, MODE_ZEROING); | ||
| 549 | + vshrs((vector_s32 &)vb32_exponent, (vector_s32 &)vb32_exponent, ctx.shr, preg_b32, MODE_ZEROING); | ||
| 550 | + vsub((vector_u32 &)vb32_shared_exp, (vector_u32 &)vb32_exponent, (vector_u32 &)ctx.vb32_b8_emax, preg_b32); | ||
| 551 | + vsub((vector_s32 &)vb32_scaling, (vector_s32 &)ctx.vb32_exp_max, (vector_s32 &)vb32_shared_exp, preg_b32); | ||
| 552 | + vshls((vector_u32 &)vb32_scaling, (vector_u32 &)vb32_scaling, ctx.shr, preg_b32, MODE_ZEROING); | ||
| 553 | + vcmps_le(ctx.preg_min_scale, (vector_s32 &)vb32_exponent, 8, preg_b32); | ||
| 554 | + vsel(vb32_scaling, ctx.vb32_recip_min_scale, vb32_scaling, ctx.preg_min_scale); | ||
| 555 | + vsel(vb32_shared_exp, ctx.vb32_zero, vb32_shared_exp, ctx.preg_min_scale); | ||
| 556 | + vcmps_eq(ctx.preg_special, (vector_s32 &)vb32_exponent, 0xFF, preg_b32); | ||
| 557 | + vcmps_ne(ctx.preg_nan, (vector_s32 &)vb32_mantissa, 0, ctx.preg_special); | ||
| 558 | + vsel(vb32_scaling, ctx.vb32_f32_nan, vb32_scaling, ctx.preg_nan); | ||
| 559 | + vsel(vb32_shared_exp, ctx.vb32_b8_nan, vb32_shared_exp, ctx.preg_nan); | ||
| 560 | +} | ||
| 561 | + | ||
| 562 | +// Shared core for fp32 OCP exponent+scaling extraction. Loads one VL from maxPtr | ||
| 563 | +// at the given element offset, computes shared exponent and reciprocal scaling, | ||
| 564 | +// and stores exponent (PK4_B32) and scaling (NORM_B32). Used by both the | ||
| 565 | +// loop-based ExtractB8ExponentAndScaling<float> and the DN-mode VL variant. | ||
| 566 | +PTO_INTERNAL void ExtractF32OcpExponentAndScalingCore(__ubuf__ float *maxPtr, __ubuf__ uint8_t *expPtr, | ||
| 567 | + __ubuf__ float *scalingPtr, uint32_t off, uint32_t elemCount) | ||
| 568 | +{ | ||
| 569 | + static constexpr auto distValue = | ||
| 570 | + std::integral_constant<::DistVST, static_cast<::DistVST>(GetDistVst<float, DistVST::DIST_NORM>())>(); | ||
| 571 | + F32OcpQuantCtx ctx; | ||
| 572 | + InitF32OcpQuantCtx(ctx); | ||
| 573 | + uint32_t preg_cols_b32 = elemCount; | ||
| 574 | + uint32_t preg_cols_b8 = elemCount * 4; | ||
| 575 | + MaskReg preg_b32 = CreatePredicate<float>(preg_cols_b32); | ||
| 576 | + MaskReg preg_b8 = CreatePredicate<uint8_t>(preg_cols_b8); | ||
| 577 | + vector_f32 vb32_max; | ||
| 578 | + vector_s32 vb32_shared_exp, vb32_scaling; | ||
| 579 | + ComputeF32OcpExpAndScaling(ctx, vb32_shared_exp, vb32_scaling, vb32_max, preg_b32, (__ubuf__ int32_t *)maxPtr, off); | ||
| 580 | + vsts((vector_s32 &)vb32_shared_exp, ((__ubuf__ int32_t *)expPtr), off / 4, PK4_B32, preg_b8); | ||
| 581 | + vsts((vector_s32 &)vb32_scaling, ((__ubuf__ int32_t *)scalingPtr), off, distValue, preg_b32); | ||
| 582 | +} | ||
| 583 | + | ||
| 584 | +// Unrolled variant: interleaves scaling stores for higher throughput. | ||
| 585 | +PTO_INTERNAL void ExtractB8ExponentAndScalingUnrolled(__ubuf__ float *maxPtr, __ubuf__ uint8_t *expPtr, | ||
| 586 | + __ubuf__ float *scalingPtr, unsigned exp_max_loop_count, | ||
| 587 | + unsigned total_elements_count, unsigned elementsPerRepeat) | ||
| 588 | +{ | ||
| 589 | + F32OcpQuantCtx ctx; | ||
| 590 | + InitF32OcpQuantCtx(ctx); | ||
| 591 | + uint32_t total_count = total_elements_count; | ||
| 592 | + uint32_t scaling_elem_count = total_elements_count * 2; | ||
| 593 | + vector_f32 vb32_max; | ||
| 594 | + vector_s32 vb32_shared_exp, vb32_scaling; | ||
| 595 | + for (uint16_t i = 0; i < (uint16_t)exp_max_loop_count; ++i) { | ||
| 596 | + MaskReg preg_b32 = CreatePredicate<float>(total_count); | ||
| 597 | + ComputeF32OcpExpAndScaling(ctx, vb32_shared_exp, vb32_scaling, vb32_max, preg_b32, (__ubuf__ int32_t *)maxPtr, | ||
| 598 | + i * elementsPerRepeat); | ||
| 599 | + vsts((vector_s32 &)vb32_shared_exp, ((__ubuf__ int32_t *)expPtr), i * elementsPerRepeat / 4, PK4_B32, preg_b32); | ||
| 600 | + vector_s32 vb32_scaling_0, vb32_scaling_1; | ||
| 601 | + vintlv(vb32_scaling_0, vb32_scaling_1, vb32_scaling, vb32_scaling); | ||
| 602 | + MaskReg preg_sc0 = CreatePredicate<float>(scaling_elem_count); | ||
| 603 | + MaskReg preg_sc1 = CreatePredicate<float>(scaling_elem_count); | ||
| 604 | + vsts((vector_s32 &)vb32_scaling_0, ((__ubuf__ int32_t *)scalingPtr), 2 * i * elementsPerRepeat, NORM_B32, | ||
| 605 | + preg_sc0); | ||
| 606 | + vsts((vector_s32 &)vb32_scaling_1, ((__ubuf__ int32_t *)scalingPtr + 64), 2 * i * elementsPerRepeat, NORM_B32, | ||
| 607 | + preg_sc1); | ||
| 608 | + } | ||
| 609 | +} | ||
| 610 | + | ||
| 611 | +// Assumptions: Input is float, data is continuous and 1D, and the usual assumptions about M (divisible by 64) | ||
| 518 | // Computing scalar focus and exponent for F32 -> b8 e4m3 quantization | 612 | // Computing scalar focus and exponent for F32 -> b8 e4m3 quantization |
| 519 | template <bool unroll = false> | 613 | template <bool unroll = false> |
| 520 | PTO_INTERNAL void ExtractB8ExponentAndScaling(__ubuf__ float *maxPtr, __ubuf__ uint8_t *expPtr, | 614 | PTO_INTERNAL void ExtractB8ExponentAndScaling(__ubuf__ float *maxPtr, __ubuf__ uint8_t *expPtr, |
| 521 | __ubuf__ float *scalingPtr, unsigned exp_max_loop_count, | 615 | __ubuf__ float *scalingPtr, unsigned exp_max_loop_count, |
| 522 | unsigned total_elements_count, unsigned elementsPerRepeat) | 616 | unsigned total_elements_count, unsigned elementsPerRepeat) |
| 523 | { | 617 | { |
| 524 | - static constexpr auto distValue = | 618 | + if constexpr (unroll) { |
| 525 | - std::integral_constant<::DistVST, static_cast<::DistVST>(GetDistVst<float, DistVST::DIST_NORM>())>(); | 619 | + ExtractB8ExponentAndScalingUnrolled(maxPtr, expPtr, scalingPtr, exp_max_loop_count, total_elements_count, |
| 526 | - vector_f32 vb32_max; | 620 | + elementsPerRepeat); |
| 527 | - vector_s32 vb32_exponent, vb32_mantissa, vb32_shared_exp, vb32_scaling; | 621 | + } else { |
| 528 | - vector_s32 vb32_b8_nan, vb32_f32_nan, vb32_b8_emax, vb32_exp_mask, vb32_mantissa_mask, vb32_exp_max; | 622 | + for (uint16_t i = 0; i < (uint16_t)exp_max_loop_count; ++i) { |
| 529 | - vector_s32 vb32_recip_min_scale, vb32_zero; | 623 | + ExtractF32OcpExponentAndScalingCore(maxPtr, expPtr, scalingPtr, i * elementsPerRepeat, |
| 530 | - constexpr int shr = 23; | 624 | + total_elements_count); |
| 531 | - vbr(vb32_exp_mask, 0x7F800000); | 625 | + } |
| 532 | - vbr(vb32_mantissa_mask, 0x007FFFFF); | ||
| 533 | - vbr(vb32_b8_nan, 0xFF); | ||
| 534 | - vbr(vb32_f32_nan, 0x7FC00000); | ||
| 535 | - vbr(vb32_exp_max, 0xFE); | ||
| 536 | - vbr(vb32_b8_emax, 8); // Max exponent for e4m3 is 8 | ||
| 537 | - vbr(vb32_recip_min_scale, 0x7F000000); | ||
| 538 | - vbr(vb32_zero, 0); | ||
| 539 | - vector_bool preg_special, preg_nan, preg_min_scale; | ||
| 540 | - uint32_t total_count = total_elements_count; | ||
| 541 | - uint32_t scaling_elem_count = total_elements_count * 2; | ||
| 542 | - for (uint16_t i = 0; i < (uint16_t)exp_max_loop_count; ++i) { | ||
| 543 | - vector_bool preg_b32 = CreatePredicate<float>(total_count); | ||
| 544 | - vlds((vector_s32 &)vb32_max, (__ubuf__ int32_t *)maxPtr, i * elementsPerRepeat, NORM); | ||
| 545 | - vand((vector_s32 &)vb32_exponent, (vector_s32 &)vb32_max, vb32_exp_mask, preg_b32, MODE_ZEROING); | ||
| 546 | - vand((vector_s32 &)vb32_mantissa, (vector_s32 &)vb32_max, vb32_mantissa_mask, preg_b32, MODE_ZEROING); | ||
| 547 | - vshrs((vector_s32 &)vb32_exponent, (vector_s32 &)vb32_exponent, shr, preg_b32, MODE_ZEROING); | ||
| 548 | - vsub((vector_u32 &)vb32_shared_exp, (vector_u32 &)vb32_exponent, (vector_u32 &)vb32_b8_emax, preg_b32); | ||
| 549 | - vsub((vector_s32 &)vb32_scaling, (vector_s32 &)vb32_exp_max, (vector_s32 &)vb32_shared_exp, preg_b32); | ||
| 550 | - vshls((vector_u32 &)vb32_scaling, (vector_u32 &)vb32_scaling, shr, preg_b32, MODE_ZEROING); | ||
| 551 | - | ||
| 552 | - vcmps_le(preg_min_scale, (vector_s32 &)vb32_exponent, 8, preg_b32); | ||
| 553 | - vsel(vb32_scaling, vb32_recip_min_scale, vb32_scaling, preg_min_scale); | ||
| 554 | - vsel(vb32_shared_exp, vb32_zero, vb32_shared_exp, preg_min_scale); | ||
| 555 | - | ||
| 556 | - vcmps_eq(preg_special, (vector_s32 &)vb32_exponent, 0xFF, preg_b32); | ||
| 557 | - vcmps_ne(preg_nan, (vector_s32 &)vb32_mantissa, 0, preg_special); | ||
| 558 | - vsel(vb32_scaling, vb32_f32_nan, vb32_scaling, preg_nan); | ||
| 559 | - vsel(vb32_shared_exp, vb32_b8_nan, vb32_shared_exp, preg_nan); | ||
| 560 | - vsts((vector_s32 &)vb32_shared_exp, ((__ubuf__ int32_t *)expPtr), i * elementsPerRepeat / 4, PK4_B32, preg_b32); | ||
| 561 | - if constexpr (unroll) { | ||
| 562 | - vector_s32 vb32_scaling_0, vb32_scaling_1; | ||
| 563 | - vintlv(vb32_scaling_0, vb32_scaling_1, vb32_scaling, vb32_scaling); | ||
| 564 | - MaskReg preg_scaling_0 = CreatePredicate<float>(scaling_elem_count); | ||
| 565 | - MaskReg preg_scaling_1 = CreatePredicate<float>(scaling_elem_count); | ||
| 566 | - vsts((vector_s32 &)vb32_scaling_0, ((__ubuf__ int32_t *)scalingPtr), 2 * i * elementsPerRepeat, NORM_B32, | ||
| 567 | - preg_scaling_0); | ||
| 568 | - vsts((vector_s32 &)vb32_scaling_1, ((__ubuf__ int32_t *)scalingPtr + 64), 2 * i * elementsPerRepeat, | ||
| 569 | - NORM_B32, preg_scaling_1); | ||
| 570 | - | ||
| 571 | - } else | ||
| 572 | - vsts((vector_s32 &)vb32_scaling, ((__ubuf__ int32_t *)scalingPtr), i * elementsPerRepeat, distValue, | ||
| 573 | - preg_b32); | ||
| 574 | } | 626 | } |
| 575 | } | 627 | } |
| 576 | 628 | ||
| @@ -734,6 +786,14 @@ PTO_INTERNAL void ExtractB8ExponentAndScalingVL(__ubuf__ T *maxPtr, __ubuf__ uin | |||
| 734 | ExtractMxOcpExponentAndScalingVL<T, OcpMxFp8E4M3Spec>(maxPtr, expPtr, scalingPtr, off, rem); | 786 | ExtractMxOcpExponentAndScalingVL<T, OcpMxFp8E4M3Spec>(maxPtr, expPtr, scalingPtr, off, rem); |
| 735 | } | 787 | } |
| 736 | 788 | ||
| 789 | +// fp32 specialization of ExtractB8ExponentAndScalingVL for DN mode. | ||
| 790 | +template <> | ||
| 791 | +PTO_INTERNAL void ExtractB8ExponentAndScalingVL<float>(__ubuf__ float *maxPtr, __ubuf__ uint8_t *expPtr, | ||
| 792 | + __ubuf__ float *scalingPtr, uint32_t off, uint32_t rem) | ||
| 793 | +{ | ||
| 794 | + ExtractF32OcpExponentAndScalingCore(maxPtr, expPtr, scalingPtr, off, rem); | ||
| 795 | +} | ||
| 796 | + | ||
| 737 | template <typename T> | 797 | template <typename T> |
| 738 | PTO_INTERNAL void ExtractB8ExponentAndScaling(__ubuf__ T *maxPtr, __ubuf__ uint8_t *expPtr, __ubuf__ T *scalingPtr, | 798 | PTO_INTERNAL void ExtractB8ExponentAndScaling(__ubuf__ T *maxPtr, __ubuf__ uint8_t *expPtr, __ubuf__ T *scalingPtr, |
| 739 | unsigned exp_max_loop_count, unsigned total_elements_count) | 799 | unsigned exp_max_loop_count, unsigned total_elements_count) |
| @@ -2106,14 +2166,14 @@ __tf__ PTO_INTERNAL void TQuant_MXFP8_Impl(typename TileDataOut::TileDType __out | |||
| 2106 | typename TileDataSrc::TileDType __in__ src, unsigned validRows, | 2166 | typename TileDataSrc::TileDType __in__ src, unsigned validRows, |
| 2107 | unsigned validCols) | 2167 | unsigned validCols) |
| 2108 | { | 2168 | { |
| 2109 | - using T = typename TileDataSrc::DType; | ||
| 2110 | - using ExpT = typename TileDataExp::DType; | ||
| 2111 | using OutT = typename TileDataOut::DType; | 2169 | using OutT = typename TileDataOut::DType; |
| 2112 | - __ubuf__ T *srcPtr = (__ubuf__ T *)__cce_get_tile_ptr(src); | 2170 | + using ExpT = typename TileDataExp::DType; |
| 2171 | + using T = typename TileDataSrc::DType; | ||
| 2113 | __ubuf__ ExpT *expPtr = (__ubuf__ ExpT *)__cce_get_tile_ptr(exp); | 2172 | __ubuf__ ExpT *expPtr = (__ubuf__ ExpT *)__cce_get_tile_ptr(exp); |
| 2114 | - __ubuf__ OutT *dstPtr = (__ubuf__ OutT *)__cce_get_tile_ptr(dst); | ||
| 2115 | __ubuf__ T *maxPtr = (__ubuf__ T *)__cce_get_tile_ptr(max); | 2173 | __ubuf__ T *maxPtr = (__ubuf__ T *)__cce_get_tile_ptr(max); |
| 2174 | + __ubuf__ T *srcPtr = (__ubuf__ T *)__cce_get_tile_ptr(src); | ||
| 2116 | __ubuf__ T *scalingPtr = (__ubuf__ T *)__cce_get_tile_ptr(scaling); | 2175 | __ubuf__ T *scalingPtr = (__ubuf__ T *)__cce_get_tile_ptr(scaling); |
| 2176 | + __ubuf__ OutT *dstPtr = (__ubuf__ OutT *)__cce_get_tile_ptr(dst); | ||
| 2117 | 2177 | ||
| 2118 | set_ctrl(static_cast<uint64_t>(1) << 50); | 2178 | set_ctrl(static_cast<uint64_t>(1) << 50); |
| 2119 | __VEC_SCOPE__ | 2179 | __VEC_SCOPE__ |
| @@ -2271,7 +2331,205 @@ __tf__ PTO_INTERNAL void TQuant_Int8Asym(typename TileDataOut::TileDType __out__ | |||
| 2271 | } | 2331 | } |
| 2272 | } | 2332 | } |
| 2273 | 2333 | ||
| 2274 | -// TQuant Interface for FP32/FP16/BF16->INT4/8/16 | 2334 | +// Missing stuff and TODOS: |
| 2335 | +// 1) Dynamic vs static predicate implementation | ||
| 2336 | +// 2) Testing on board to assure this case does not fail | ||
| 2337 | +// 3) Loop peeling is more efficient than using vbr, but just get the correctness first then use loop peeling | ||
| 2338 | +// Assumptions: | ||
| 2339 | +// 1) validRows is divisible by 32 (grpSize) | ||
| 2340 | +template <typename T, uint32_t StaticCols> | ||
| 2341 | +PTO_INTERNAL void AbsReduceMax_DN(__ubuf__ T *srcPtr, __ubuf__ T *maxPtr, unsigned validRows, unsigned validCols) | ||
| 2342 | +{ | ||
| 2343 | + constexpr uint32_t grpSize = 32; | ||
| 2344 | + constexpr uint32_t elementsPerVL = REPEAT_BYTE / sizeof(T); | ||
| 2345 | + uint32_t num_vls_per_row = CeilDivision(validCols, elementsPerVL); | ||
| 2346 | + uint32_t num_grps_per_col = CeilDivision(validRows, grpSize); | ||
| 2347 | + constexpr uint32_t num_vls_inner_loop = 4; | ||
| 2348 | + uint32_t inner_loop_iters = CeilDivision(grpSize, num_vls_inner_loop); | ||
| 2349 | + static constexpr auto distValue = | ||
| 2350 | + std::integral_constant<::DistVST, static_cast<::DistVST>(GetDistVst<T, DistVST::DIST_NORM>())>(); | ||
| 2351 | + // vabs expects a floating-point vector type. bf16 registers are represented | ||
| 2352 | + // as vector_f16 on this hardware, so cast accordingly; fp32 uses vector_f32. | ||
| 2353 | + using AbsVecType = std::conditional_t<std::is_same<T, float>::value, vector_f32, vector_f16>; | ||
| 2354 | + RegTensor<T> vreg_0, vreg_1, vreg_2, vreg_3; | ||
| 2355 | + RegTensor<T> vreg_max, vreg_max_0, vreg_max_1, vreg_max_2, vreg_max_3; | ||
| 2356 | + uint32_t preg_cols = validCols; | ||
| 2357 | + for (uint32_t i = 0; i < num_vls_per_row; ++i) { | ||
| 2358 | + uint32_t vl_start = i * elementsPerVL; | ||
| 2359 | + MaskReg preg = CreatePredicate<T>(preg_cols); | ||
| 2360 | + for (uint32_t j = 0; j < num_grps_per_col; ++j) { | ||
| 2361 | + vbr(vreg_max_0, (T)0); | ||
| 2362 | + vbr(vreg_max_1, (T)0); | ||
| 2363 | + vbr(vreg_max_2, (T)0); | ||
| 2364 | + vbr(vreg_max_3, (T)0); | ||
| 2365 | + uint32_t grp_start = j * grpSize * StaticCols; | ||
| 2366 | + for (uint32_t k = 0; k < inner_loop_iters; ++k) { | ||
| 2367 | + uint32_t inner_start = k * num_vls_inner_loop * StaticCols; | ||
| 2368 | + uint32_t offset = vl_start + grp_start + inner_start; | ||
| 2369 | + vlds(vreg_0, srcPtr + offset, 0, NORM); | ||
| 2370 | + vlds(vreg_1, srcPtr + offset, 1 * StaticCols, NORM); | ||
| 2371 | + vlds(vreg_2, srcPtr + offset, 2 * StaticCols, NORM); | ||
| 2372 | + vlds(vreg_3, srcPtr + offset, 3 * StaticCols, NORM); | ||
| 2373 | + vabs((AbsVecType &)vreg_0, (AbsVecType &)vreg_0, preg); | ||
| 2374 | + vabs((AbsVecType &)vreg_1, (AbsVecType &)vreg_1, preg); | ||
| 2375 | + vabs((AbsVecType &)vreg_2, (AbsVecType &)vreg_2, preg); | ||
| 2376 | + vabs((AbsVecType &)vreg_3, (AbsVecType &)vreg_3, preg); | ||
| 2377 | + vmax(vreg_max_0, vreg_0, vreg_max_0, preg, MODE_ZEROING); | ||
| 2378 | + vmax(vreg_max_1, vreg_1, vreg_max_1, preg, MODE_ZEROING); | ||
| 2379 | + vmax(vreg_max_2, vreg_2, vreg_max_2, preg, MODE_ZEROING); | ||
| 2380 | + vmax(vreg_max_3, vreg_3, vreg_max_3, preg, MODE_ZEROING); | ||
| 2381 | + } | ||
| 2382 | + vmax(vreg_max_0, vreg_max_0, vreg_max_1, preg, MODE_ZEROING); | ||
| 2383 | + vmax(vreg_max_2, vreg_max_2, vreg_max_3, preg, MODE_ZEROING); | ||
| 2384 | + vmax(vreg_max, vreg_max_0, vreg_max_2, preg, MODE_ZEROING); | ||
| 2385 | + vsts(vreg_max, maxPtr, j * StaticCols + i * elementsPerVL, distValue, preg); | ||
| 2386 | + } | ||
| 2387 | + } | ||
| 2388 | +} | ||
| 2389 | + | ||
| 2390 | +// fp16 & bf16 | ||
| 2391 | +template <typename T, uint32_t StaticCols> | ||
| 2392 | +PTO_INTERNAL void calcQuantizedFP8Values_DN_B16(__ubuf__ T *srcPtr, __ubuf__ T *scalingPtr, __ubuf__ uint8_t *dstPtr, | ||
| 2393 | + unsigned validRows, unsigned validCols) | ||
| 2394 | +{ | ||
| 2395 | + constexpr uint32_t grpSize = 32; | ||
| 2396 | + constexpr uint32_t b16ElementsPerVL = REPEAT_BYTE / sizeof(T); // B16 elements per VL | ||
| 2397 | + uint32_t num_vls_per_row = CeilDivision((uint32_t)validCols, b16ElementsPerVL); | ||
| 2398 | + uint32_t num_grps_per_col = CeilDivision((uint32_t)validRows, grpSize); | ||
| 2399 | + RegTensor<T> vb16_scaling, vb16_input; | ||
| 2400 | + vector_f32 vb32_scaling_even, vb32_scaling_odd; | ||
| 2401 | + vector_f32 vb32_input_even, vb32_input_odd; | ||
| 2402 | + vector_f8e4m3 vb8_p0, vb8_p1, vb8_out; | ||
| 2403 | + uint32_t preg_cols_b16 = validCols; | ||
| 2404 | + uint32_t preg_cols_b8 = validCols * 2; | ||
| 2405 | + uint32_t preg_cols_b32 = validCols; | ||
| 2406 | + for (uint32_t i = 0; i < num_vls_per_row; ++i) { | ||
| 2407 | + uint32_t vl_start = i * b16ElementsPerVL; | ||
| 2408 | + MaskReg preg_b16 = CreatePredicate<bfloat16_t>(preg_cols_b16); | ||
| 2409 | + MaskReg preg_b8 = CreatePredicate<uint8_t>(preg_cols_b8); | ||
| 2410 | + MaskReg preg_b32 = CreatePredicate<float>(preg_cols_b32); | ||
| 2411 | + for (uint32_t j = 0; j < num_grps_per_col; ++j) { | ||
| 2412 | + uint32_t row_base = j * grpSize; | ||
| 2413 | + vlds(vb16_scaling, scalingPtr, vl_start + j * StaticCols, NORM); | ||
| 2414 | + vcvt(vb32_scaling_even, vb16_scaling, preg_b16, PART_EVEN); | ||
| 2415 | + vcvt(vb32_scaling_odd, vb16_scaling, preg_b16, PART_ODD); | ||
| 2416 | + for (uint32_t k = 0; k < grpSize; ++k) { | ||
| 2417 | + uint32_t r = row_base + k; | ||
| 2418 | + vlds(vb16_input, srcPtr, r * StaticCols + vl_start, NORM); | ||
| 2419 | + vcvt(vb32_input_even, vb16_input, preg_b16, PART_EVEN); | ||
| 2420 | + vcvt(vb32_input_odd, vb16_input, preg_b16, PART_ODD); | ||
| 2421 | + vmul(vb32_input_even, vb32_input_even, vb32_scaling_even, preg_b32, MODE_ZEROING); | ||
| 2422 | + vmul(vb32_input_odd, vb32_input_odd, vb32_scaling_odd, preg_b32, MODE_ZEROING); | ||
| 2423 | + vcvt(vb8_p0, vb32_input_even, preg_b32, ROUND_R, RS_ENABLE, PART_P0); | ||
| 2424 | + vcvt(vb8_p1, vb32_input_odd, preg_b32, ROUND_R, RS_ENABLE, PART_P1); | ||
| 2425 | + vor(vb8_out, vb8_p0, vb8_p1, preg_b8); | ||
| 2426 | + uint32_t dst_byte_offset = r * StaticCols + vl_start; | ||
| 2427 | + vsts((vector_u32 &)vb8_out, (__ubuf__ uint32_t *)dstPtr, dst_byte_offset / 4, PK_B32, preg_b8); | ||
🟡 Medium Priority 变更位置:第 2411 行 在整个 pre-existing 代码库中,PK_B32 存储模式搭配的是 32-bit 谓词( 此处新代码使用 注意: 建议:将第 2411 行的 preg_b8 改为 preg_b32(该变量已在第 2394 行创建),与 PK_B32 存储模式的 32-bit 谓词语义保持一致。 ![]() ![]() 不准确? | |||
| 2428 | + } | ||
| 2429 | + } | ||
| 2430 | + } | ||
| 2431 | +} | ||
| 2432 | + | ||
| 2433 | +// fp32 | ||
| 2434 | +template <uint32_t StaticCols> | ||
| 2435 | +PTO_INTERNAL void calcQuantizedFP8Values_DN_float(__ubuf__ float *srcPtr, __ubuf__ float *scalingPtr, | ||
| 2436 | + __ubuf__ uint8_t *dstPtr, unsigned validRows, unsigned validCols) | ||
| 2437 | +{ | ||
| 2438 | + constexpr uint32_t grpSize = 32; | ||
| 2439 | + constexpr uint32_t b32ElementsPerVL = REPEAT_BYTE / sizeof(float); // B32 elements per VL | ||
| 2440 | + uint32_t num_vls_per_row = CeilDivision((uint32_t)validCols, b32ElementsPerVL); | ||
| 2441 | + uint32_t num_grps_per_col = CeilDivision((uint32_t)validRows, grpSize); | ||
| 2442 | + RegTensor<float> vf32_scaling, vf32_input; | ||
| 2443 | + vector_f8e4m3 vb8_out; | ||
| 2444 | + for (uint32_t i = 0; i < num_vls_per_row; ++i) { | ||
| 2445 | + uint32_t vl_start = i * b32ElementsPerVL; | ||
| 2446 | + uint32_t preg_cols_b32 = validCols; | ||
| 2447 | + uint32_t preg_cols_b8 = validCols * 4; | ||
| 2448 | + MaskReg preg_b32 = CreatePredicate<float>(preg_cols_b32); | ||
| 2449 | + MaskReg preg_b8 = CreatePredicate<uint8_t>(preg_cols_b8); | ||
| 2450 | + for (uint32_t j = 0; j < num_grps_per_col; ++j) { | ||
| 2451 | + uint32_t row_base = j * grpSize; | ||
| 2452 | + vlds(vf32_scaling, scalingPtr, vl_start + j * StaticCols, NORM); | ||
| 2453 | + for (uint32_t k = 0; k < grpSize; ++k) { | ||
| 2454 | + uint32_t r = row_base + k; | ||
| 2455 | + vlds(vf32_input, srcPtr, r * StaticCols + vl_start, NORM); | ||
| 2456 | + vmul(vf32_input, vf32_input, vf32_scaling, preg_b32, MODE_ZEROING); | ||
| 2457 | + vcvt(vb8_out, vf32_input, preg_b32, ROUND_R, RS_ENABLE, PART_P0); | ||
| 2458 | + vsts((vector_u8 &)vb8_out, dstPtr, r * StaticCols + vl_start, PK4_B32, preg_b8); | ||
🟡 Medium Priority 变更位置:第 2442 行 与 Finding #1 和 #2 一致的谓词类型错误:PK4_B32 在整个 pre-existing 代码库中搭配的是 32-bit 谓词( 建议:将第 2442 行的 preg_b8 改为 preg_b32,移除不再需要的 preg_cols_b8(第 2431 行)和 preg_b8(第 2433 行)。 ![]() ![]() 不准确? | |||
| 2459 | + } | ||
| 2460 | + } | ||
| 2461 | + } | ||
| 2462 | +} | ||
| 2463 | + | ||
| 2464 | +template <QuantScaleAlg scale_alg, typename T, unsigned StaticCols> | ||
| 2465 | +PTO_INTERNAL void TQuant_MXFP8_DN(__ubuf__ T *srcPtr, __ubuf__ uint8_t *expPtr, __ubuf__ uint8_t *dstPtr, | ||
| 2466 | + __ubuf__ T *maxPtr, __ubuf__ T *scalingPtr, unsigned validRows, unsigned validCols) | ||
| 2467 | +{ | ||
| 2468 | + constexpr uint32_t grpSize = 32; | ||
| 2469 | + constexpr uint32_t elementsPerVL = REPEAT_BYTE / sizeof(T); | ||
| 2470 | + AbsReduceMax_DN<T, StaticCols>(srcPtr, maxPtr, validRows, validCols); | ||
| 2471 | + mem_bar(VST_VLD); | ||
| 2472 | + // DN-aware 2D extraction: each row-group j owns one row of width StaticCols | ||
| 2473 | + // in the max/scaling/exp tiles. Process each row as a 1D VL sequence so all | ||
| 2474 | + // offsets stay aligned while preserving the 2D tile shape. | ||
| 2475 | + uint32_t num_grps_per_col = CeilDivision((uint32_t)validRows, grpSize); | ||
| 2476 | + uint32_t num_vls_per_row = CeilDivision((uint32_t)validCols, elementsPerVL); | ||
| 2477 | + for (uint32_t j = 0; j < num_grps_per_col; ++j) { | ||
| 2478 | + __ubuf__ T *maxRowPtr = maxPtr + j * StaticCols; | ||
| 2479 | + __ubuf__ uint8_t *expRowPtr = expPtr + j * StaticCols; | ||
| 2480 | + __ubuf__ T *scalingRowPtr = scalingPtr + j * StaticCols; | ||
| 2481 | + for (uint32_t i = 0; i < num_vls_per_row; ++i) { | ||
| 2482 | + uint32_t off = i * elementsPerVL; | ||
| 2483 | + uint32_t rem = (validCols > off) ? (validCols - off) : 0; | ||
| 2484 | + if (rem > elementsPerVL) | ||
| 2485 | + rem = elementsPerVL; | ||
| 2486 | + ExtractB8ExponentAndScalingVL<T>(maxRowPtr, expRowPtr, scalingRowPtr, off, rem); | ||
| 2487 | + } | ||
| 2488 | + } | ||
| 2489 | + mem_bar(VST_VLD); | ||
| 2490 | + if constexpr (std::is_same<T, float>::value) | ||
| 2491 | + calcQuantizedFP8Values_DN_float<StaticCols>(srcPtr, scalingPtr, dstPtr, validRows, validCols); | ||
| 2492 | + else | ||
| 2493 | + calcQuantizedFP8Values_DN_B16<T, StaticCols>(srcPtr, scalingPtr, dstPtr, validRows, validCols); | ||
| 2494 | +} | ||
| 2495 | + | ||
| 2496 | +template <QuantScaleAlg scale_alg, typename TileDataOut, typename TileDataSrc, typename TileDataExp, | ||
| 2497 | + typename TileDataMax, typename TileDataScaling, typename TileDataExpDn> | ||
| 2498 | +__tf__ PTO_INTERNAL void TQuant_MXFP8_Impl_DN(typename TileDataOut::TileDType __out__ dst, | ||
| 2499 | + typename TileDataExp::TileDType __out__ exp, | ||
| 2500 | + typename TileDataMax::TileDType __out__ max, | ||
| 2501 | + typename TileDataScaling::TileDType __out__ scaling, | ||
| 2502 | + typename TileDataExpDn::TileDType __out__ expDn, | ||
| 2503 | + typename TileDataSrc::TileDType __in__ src, unsigned validRows, | ||
| 2504 | + unsigned validCols) | ||
| 2505 | +{ | ||
| 2506 | + using T = typename TileDataSrc::DType; | ||
| 2507 | + using ExpT = typename TileDataExp::DType; | ||
| 2508 | + using OutT = typename TileDataOut::DType; | ||
| 2509 | + __ubuf__ T *srcPtr = (__ubuf__ T *)__cce_get_tile_ptr(src); | ||
| 2510 | + __ubuf__ ExpT *expPtr = (__ubuf__ ExpT *)__cce_get_tile_ptr(exp); | ||
| 2511 | + __ubuf__ OutT *dstPtr = (__ubuf__ OutT *)__cce_get_tile_ptr(dst); | ||
| 2512 | + __ubuf__ T *maxPtr = (__ubuf__ T *)__cce_get_tile_ptr(max); | ||
| 2513 | + __ubuf__ T *scalingPtr = (__ubuf__ T *)__cce_get_tile_ptr(scaling); | ||
| 2514 | + __ubuf__ uint8_t *expDnPtr = (__ubuf__ uint8_t *)__cce_get_tile_ptr(expDn); | ||
🔵 Low Priority 变更位置:第 2498 行
虽然调用方(
修复:删除该行,或如果 DN exp 输出未来需要独立计算逻辑,添加 TODO 注释说明。 建议:移除未使用的 expDnPtr 声明,或添加注释说明为何声明但未使用。同时考虑在函数注释中说明 expDn 参数在此阶段仅由调用方通过 TMOV 填充。 ![]() ![]() 不准确? | |||
| 2515 | + | ||
| 2516 | + set_ctrl(static_cast<uint64_t>(1) << 50); | ||
| 2517 | + __VEC_SCOPE__ | ||
| 2518 | + { | ||
| 2519 | + ZeroPadSourceTile<T, TileDataSrc::Cols>(srcPtr, validRows, validCols); | ||
| 2520 | + mem_bar(VST_VLD); | ||
| 2521 | + if constexpr (std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value || | ||
| 2522 | + std::is_same<T, float>::value) { | ||
| 2523 | + TQuant_MXFP8_DN<scale_alg, T, TileDataSrc::Cols>(srcPtr, (__ubuf__ uint8_t *)expPtr, | ||
| 2524 | + (__ubuf__ uint8_t *)dstPtr, maxPtr, scalingPtr, validRows, | ||
| 2525 | + validCols); | ||
| 2526 | + } else { | ||
| 2527 | + static_assert( | ||
| 2528 | + std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value || std::is_same<T, float>::value, | ||
| 2529 | + "Fix: MXFP8 DN mode currently supports bf16/fp16/fp32 source only."); | ||
| 2530 | + } | ||
| 2531 | + } | ||
| 2532 | +} | ||
| 2275 | template <QuantType quant_type, typename TileDataOut, typename TileDataSrc, typename TileDataPara> | 2533 | template <QuantType quant_type, typename TileDataOut, typename TileDataSrc, typename TileDataPara> |
| 2276 | PTO_INTERNAL void TQUANT_IMPL(TileDataOut &dst, TileDataSrc &src, TileDataPara &scale, TileDataPara *offset = nullptr) | 2534 | PTO_INTERNAL void TQUANT_IMPL(TileDataOut &dst, TileDataSrc &src, TileDataPara &scale, TileDataPara *offset = nullptr) |
| 2277 | { | 2535 | { |
| @@ -2376,5 +2634,20 @@ PTO_INTERNAL void TQUANT_IMPL(TileDataOut &dst, TileDataSrc &src, TileDataExp *e | |||
| 2376 | TRESHAPE_IMPL(*exp, flatExp); | 2634 | TRESHAPE_IMPL(*exp, flatExp); |
| 2377 | } | 2635 | } |
| 2378 | } | 2636 | } |
| 2637 | + | ||
| 2638 | +template <QuantType quant_type, QuantScaleAlg scale_alg, typename TileDataOut, typename TileDataSrc, | ||
| 2639 | + typename TileDataExp, typename TileDataMax, typename TileDataScaling, typename TileDataExpDn> | ||
| 2640 | +PTO_INTERNAL void TQUANT_IMPL(TileDataOut &dst, TileDataSrc &src, TileDataExp *exp, TileDataMax *max, | ||
| 2641 | + TileDataScaling *scaling, TileDataExpDn *expDn) | ||
| 2642 | +{ | ||
| 2643 | + using T = typename TileDataSrc::DType; | ||
| 2644 | + static_assert(quant_type == QuantType::MXFP8, "Fix: DN mode overload supports MXFP8 only."); | ||
| 2645 | + static_assert(std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value || std::is_same<T, float>::value, | ||
| 2646 | + "Fix: MXFP8 DN input has to be bfloat16, float16 (half), or float32"); | ||
| 2647 | + | ||
| 2648 | + TQuant_MXFP8_Impl_DN<scale_alg, TileDataOut, TileDataSrc, TileDataExp, TileDataMax, TileDataScaling, TileDataExpDn>( | ||
| 2649 | + dst.data(), exp->data(), max->data(), scaling->data(), expDn->data(), src.data(), src.GetValidRow(), | ||
| 2650 | + src.GetValidCol()); | ||
| 2651 | +} | ||
| 2379 | } // namespace pto | 2652 | } // namespace pto |
| 2380 | 2653 | ||
| @@ -141,7 +141,6 @@ set(ALL_TESTCASES | |||
| 141 | tgather | 141 | tgather |
| 142 | tcvt | 142 | tcvt |
| 143 | tsub | 143 | tsub |
| 144 | - tsubreluconv | ||
| 145 | trem | 144 | trem |
| 146 | tfmod | 145 | tfmod |
| 147 | tprelu | 146 | tprelu |
| @@ -156,7 +155,6 @@ set(ALL_TESTCASES | |||
| 156 | tshr | 155 | tshr |
| 157 | tshrs | 156 | tshrs |
| 158 | tadd | 157 | tadd |
| 159 | - taddreluconv | ||
| 160 | tpartadd | 158 | tpartadd |
| 161 | tpartmul | 159 | tpartmul |
| 162 | tfillpad | 160 | tfillpad |
| @@ -249,11 +247,6 @@ set(ALL_TESTCASES | |||
| 249 | tpushpop_subtile | 247 | tpushpop_subtile |
| 250 | tinterleave | 248 | tinterleave |
| 251 | tdeinterleave | 249 | tdeinterleave |
| 252 | - tpairreducesum | ||
| 253 | - tfusedmuladd | ||
| 254 | - tmuladddst | ||
| 255 | - tsubrelu | ||
| 256 | - tfusedmuladdrelu | ||
| 257 | ) | 250 | ) |
| 258 | 251 | ||
| 259 | # Opt-in testcases: built ONLY when TEST_CASE=<name> is explicitly supplied. | 252 | # Opt-in testcases: built ONLY when TEST_CASE=<name> is explicitly supplied. |
| @@ -262,6 +255,7 @@ set(ALL_TESTCASES | |||
| 262 | set(OPT_IN_TESTCASES | 255 | set(OPT_IN_TESTCASES |
| 263 | # Device-side async L2 prefetch: correctness + focused perf comparison | 256 | # Device-side async L2 prefetch: correctness + focused perf comparison |
| 264 | tprefetch_async | 257 | tprefetch_async |
| 258 | +tquant_dn | ||
| 265 | ) | 259 | ) |
| 266 | 260 | ||
| 267 | if (AUTO_MODE) | 261 | if (AUTO_MODE) |
| @@ -290,6 +284,7 @@ if (AUTO_MODE) | |||
| 290 | tmatmul | 284 | tmatmul |
| 291 | tmatmul_mx | 285 | tmatmul_mx |
| 292 | tmov_zz | 286 | tmov_zz |
| 287 | + tquant_dn | ||
| 293 | ) | 288 | ) |
| 294 | 289 | ||
| 295 | foreach(whitelisted_case ${AUTO_MODE_WHITELIST}) | 290 | foreach(whitelisted_case ${AUTO_MODE_WHITELIST}) |
| @@ -0,0 +1,11 @@ | |||
| 1 | +# -------------------------------------------------------------------------------- | ||
| 2 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | +# CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | +# Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | +# See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | +# -------------------------------------------------------------------------------- | ||
| 10 | + | ||
| 11 | +pto_vec_st(tquant_dn) | ||
| @@ -0,0 +1,216 @@ | |||
| 1 | +#!/usr/bin/python3 | ||
| 2 | +# coding=utf-8 | ||
| 3 | +# -------------------------------------------------------------------------------- | ||
| 4 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 5 | +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 6 | +# CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 7 | +# Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 8 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 9 | +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 10 | +# See LICENSE in the root of the software repository for the full text of the License. | ||
| 11 | +# -------------------------------------------------------------------------------- | ||
| 12 | + | ||
| 13 | +import math | ||
| 14 | +import os | ||
| 15 | + | ||
| 16 | +import numpy as np | ||
| 17 | + | ||
| 18 | + | ||
| 19 | +def fp32_to_bf16_bits(x): | ||
| 20 | + x = np.asarray(x, dtype=np.float32) | ||
| 21 | + u32 = x.view(np.uint32) | ||
| 22 | + u16 = (u32 >> 16).astype(np.uint16) | ||
| 23 | + return u16 | ||
| 24 | + | ||
| 25 | + | ||
| 26 | +def bf16_bits_to_fp32(bf16_bits): | ||
| 27 | + u32 = np.array(bf16_bits, dtype=np.uint32) << 16 | ||
| 28 | + return u32.view(np.float32) | ||
| 29 | + | ||
| 30 | + | ||
| 31 | +def get_group_max_dn(src, group_size=32): | ||
| 32 | + m, n = src.shape | ||
| 33 | + hat_m = m // group_size | ||
| 34 | + max_vals = np.zeros((hat_m, n), dtype=np.float32) | ||
| 35 | + for rb in range(hat_m): | ||
| 36 | + for c in range(n): | ||
| 37 | + max_vals[rb, c] = np.max(np.abs(src[rb * group_size : (rb + 1) * group_size, c])) | ||
| 38 | + return max_vals | ||
| 39 | + | ||
| 40 | + | ||
| 41 | +def fp32_maxes_to_fp8(group_max, emax=8): | ||
| 42 | + max_bits = np.asarray(group_max, dtype=np.float32).view(np.uint32) | ||
| 43 | + exponent_b32 = (max_bits & 0x7F800000) >> 23 | ||
| 44 | + e8m0 = exponent_b32.astype(np.int32) - emax | ||
| 45 | + e8m0 = np.clip(e8m0, 0, 254).astype(np.uint8) | ||
| 46 | + scale_exp = 254 - e8m0.astype(np.int32) | ||
| 47 | + scale_exp = np.clip(scale_exp, 0, 255).astype(np.uint32) | ||
| 48 | + scaling_bits = (scale_exp << 23).view(np.float32) | ||
| 49 | + nan_mask = exponent_b32 == 255 | ||
| 50 | + e8m0[nan_mask] = 0xFF | ||
| 51 | + scaling_bits[nan_mask] = np.float32(np.nan) | ||
| 52 | + return e8m0, scaling_bits | ||
🟡 Medium Priority 变更位置: Python golden 数据生成器将所有 exponent == 255 的值都标记为 NaN(第 46 行 但内核
Inf(exponent=255, mantissa=0)在 kernel 中走正常路径(shared_exp=247, scaling=2^-120),但 Python 中却被写入 NaN scaling 和 0xFF exponent。这使得 golden reference 在涉及 Inf 输入时与内核输出不一致。 当前随机测试数据未产生 Inf/NaN,故不影响当前测试结果。但当生产数据中出现 Inf 时,golden 对比将错误报告 mismatch。 修复方向:将 Python 中的 建议:将 nan_mask 改为同时检查 mantissa != 0: ![]() ![]() 不准确? | |||
| 53 | + | ||
| 54 | + | ||
| 55 | +def scale_data_dn(src, scaling, group_size=32): | ||
| 56 | + m, n = src.shape | ||
| 57 | + hat_m = m // group_size | ||
| 58 | + result = np.zeros_like(src) | ||
| 59 | + for rb in range(hat_m): | ||
| 60 | + for r in range(rb * group_size, (rb + 1) * group_size): | ||
| 61 | + result[r, :] = src[r, :] * scaling[rb, :] | ||
| 62 | + return result | ||
| 63 | + | ||
| 64 | + | ||
| 65 | +def fp32_to_e4m3(x): | ||
| 66 | + from ml_dtypes import float8_e4m3fn | ||
| 67 | + | ||
| 68 | + x = np.asarray(x, dtype=np.float32) | ||
| 69 | + clipped = np.clip(x, -448.0, 448.0) | ||
| 70 | + result = clipped.astype(float8_e4m3fn) | ||
| 71 | + return result.view(np.uint8) | ||
| 72 | + | ||
| 73 | + | ||
| 74 | +def nd2nz_mxfp8(data_fp8, m, n): | ||
| 75 | + padded_rows16 = ((m + 15) // 16) * 16 | ||
| 76 | + virtual_row = padded_rows16 + 1 | ||
| 77 | + padded_cols = ((n + 31) // 32) * 32 | ||
| 78 | + n_col_groups = padded_cols // 32 | ||
| 79 | + nz = np.zeros(virtual_row * padded_cols, dtype=np.int8) | ||
| 80 | + data_flat = data_fp8.reshape(-1) if data_fp8.ndim > 1 else data_fp8 | ||
| 81 | + for cg in range(n_col_groups): | ||
| 82 | + for r in range(padded_rows16): | ||
| 83 | + src_idx = r * padded_cols + cg * 32 | ||
| 84 | + dst_idx = cg * virtual_row * 32 + r * 32 | ||
| 85 | + if r < m: | ||
| 86 | + nz[dst_idx : dst_idx + 32] = data_flat[src_idx : src_idx + 32] | ||
| 87 | + else: | ||
| 88 | + nz[dst_idx : dst_idx + 32] = 0 | ||
| 89 | + return nz | ||
| 90 | + | ||
| 91 | + | ||
| 92 | +def pack_e8_dn(e8m0, hat_m, n, padded_cols): | ||
| 93 | + """Row-major E8M0 tile (hat_m x padded_cols).""" | ||
| 94 | + e8m0_dn = np.zeros(hat_m * padded_cols, dtype=np.uint8) | ||
| 95 | + for rb in range(hat_m): | ||
| 96 | + for c in range(n): | ||
| 97 | + e8m0_dn[rb * padded_cols + c] = e8m0[rb, c] | ||
| 98 | + return e8m0_dn | ||
| 99 | + | ||
| 100 | + | ||
| 101 | +def dn2zz_e8m0(e8m0_dn, hat_m, n): | ||
| 102 | + # Row-major (ND) input -> ZZ is currently a flattened identity in this ST. | ||
| 103 | + return e8m0_dn[: hat_m * n].copy() | ||
| 104 | + | ||
| 105 | + | ||
| 106 | +def quant_bf16_to_mxfp8_dn(src_bf16_fp32, m, n_pad): | ||
| 107 | + src_fp32 = src_bf16_fp32 | ||
| 108 | + padded_cols = int(math.ceil(n_pad / 32) * 32) | ||
| 109 | + hat_m = m // 32 | ||
| 110 | + num_groups_flat = m * (padded_cols // 32) | ||
| 111 | + num_groups_flat_aligned = int(math.ceil(num_groups_flat / 32) * 32) | ||
| 112 | + | ||
| 113 | + group_max = get_group_max_dn(src_fp32, group_size=32) | ||
| 114 | + e8m0, scaling = fp32_maxes_to_fp8(group_max) | ||
| 115 | + scaled = scale_data_dn(src_fp32, scaling, group_size=32) | ||
| 116 | + fp8 = fp32_to_e4m3(scaled).reshape(m, n_pad) | ||
| 117 | + | ||
| 118 | + fp8_padded = np.zeros((m, padded_cols), dtype=np.int8) | ||
| 119 | + fp8_padded[:, :n_pad] = fp8 | ||
| 120 | + fp8_nd = fp8_padded.reshape(-1) | ||
| 121 | + fp8_nz = nd2nz_mxfp8(fp8_padded, m, n_pad) | ||
| 122 | + | ||
| 123 | + e8_dn = pack_e8_dn(e8m0, hat_m, n_pad, padded_cols) | ||
| 124 | + e8_zz = dn2zz_e8m0(e8_dn, hat_m, n_pad) | ||
| 125 | + if e8_zz.size < num_groups_flat_aligned: | ||
| 126 | + e8_zz_padded = np.zeros(num_groups_flat_aligned, dtype=np.uint8) | ||
| 127 | + e8_zz_padded[: e8_zz.size] = e8_zz | ||
| 128 | + e8_zz = e8_zz_padded | ||
| 129 | + elif e8_zz.size > num_groups_flat_aligned: | ||
| 130 | + e8_zz = e8_zz[:num_groups_flat_aligned] | ||
| 131 | + | ||
| 132 | + return fp8_nd, e8_dn, fp8_nz, e8_zz | ||
| 133 | + | ||
| 134 | + | ||
| 135 | +CASE_PARAMS = [ | ||
| 136 | + ("TQUANTDNTest.case_bf16_64x64", 64, 64), | ||
| 137 | + ("TQUANTDNTest.case_bf16_128x64", 128, 64), | ||
| 138 | + ("TQUANTDNTest.case_bf16_64x128", 64, 128), | ||
| 139 | + ("TQUANTDNTest.case_bf16_128x128", 128, 128), | ||
| 140 | + ("TQUANTDNTest.case_bf16_64x256", 64, 256), | ||
| 141 | + ("TQUANTDNTest.case_bf16_128x256", 128, 256), | ||
| 142 | + ("TQUANTDNTest.case_bf16_256x64", 256, 64), | ||
| 143 | + ("TQUANTDNTest.case_bf16_256x128", 256, 128), | ||
| 144 | +] | ||
| 145 | + | ||
| 146 | +FP32_CASE_PARAMS = [ | ||
| 147 | + ("TQUANTDNTest.case_fp32_64x128", 64, 128), | ||
| 148 | + ("TQUANTDNTest.case_fp32_128x128", 128, 128), | ||
| 149 | + ("TQUANTDNTest.case_fp32_64x256", 64, 256), | ||
| 150 | +] | ||
| 151 | + | ||
| 152 | +GOLDEN_DIR = os.environ.get("PTO_GOLDEN_DIR", ".") | ||
| 153 | + | ||
| 154 | + | ||
| 155 | +def _gen_src(m, n_pad): | ||
| 156 | + """Generate source data with log-uniform per-group max in [0.25, 16] * 10000.""" | ||
| 157 | + hat_m = m // 32 | ||
| 158 | + log_min = np.log2(0.25) | ||
| 159 | + log_max = np.log2(16.0) | ||
| 160 | + log_group_max = np.random.uniform(log_min, log_max, size=(hat_m, n_pad)) | ||
| 161 | + group_max_target = (2.0**log_group_max).astype(np.float32) | ||
| 162 | + base = np.random.uniform(0.1, 1.0, size=(m, n_pad)).astype(np.float32) | ||
| 163 | + group_max_repeated = np.repeat(group_max_target, 32, axis=0)[:m, :] | ||
| 164 | + return base * group_max_repeated * 10000.0 | ||
| 165 | + | ||
| 166 | + | ||
| 167 | +def _write_golden(out_dir, input_bytes, fp8_nd, e8_dn, group_max_bytes): | ||
| 168 | + os.makedirs(out_dir, exist_ok=True) | ||
| 169 | + with open(os.path.join(out_dir, "input.bin"), "wb") as f: | ||
| 170 | + f.write(input_bytes) | ||
| 171 | + with open(os.path.join(out_dir, "golden_fp8_nd.bin"), "wb") as f: | ||
| 172 | + f.write(fp8_nd.tobytes()) | ||
| 173 | + with open(os.path.join(out_dir, "golden_e8_dn.bin"), "wb") as f: | ||
| 174 | + f.write(e8_dn.tobytes()) | ||
| 175 | + with open(os.path.join(out_dir, "golden_group_max.bin"), "wb") as f: | ||
| 176 | + f.write(group_max_bytes) | ||
| 177 | + | ||
| 178 | + | ||
| 179 | +def gen_golden_data(case_name, m, n): | ||
| 180 | + n_pad = n | ||
| 181 | + src = _gen_src(m, n_pad) | ||
| 182 | + bf16_bits = fp32_to_bf16_bits(src).reshape(m, n_pad) | ||
| 183 | + src_bf16_fp32 = bf16_bits_to_fp32(bf16_bits.flatten()).reshape(m, n_pad) | ||
| 184 | + | ||
| 185 | + fp8_nd, e8_dn, _, _ = quant_bf16_to_mxfp8_dn(src_bf16_fp32, m, n_pad) | ||
| 186 | + | ||
| 187 | + group_max = get_group_max_dn(src_bf16_fp32, group_size=32) | ||
| 188 | + golden_group_max_bf16 = fp32_to_bf16_bits(group_max) | ||
| 189 | + | ||
| 190 | + out_dir = os.path.join(GOLDEN_DIR, case_name) | ||
| 191 | + _write_golden(out_dir, bf16_bits.reshape(-1).tobytes(), fp8_nd, e8_dn, golden_group_max_bf16.reshape(-1).tobytes()) | ||
| 192 | + | ||
| 193 | + | ||
| 194 | +def gen_golden_data_fp32(case_name, m, n): | ||
| 195 | + n_pad = n | ||
| 196 | + src = _gen_src(m, n_pad) | ||
| 197 | + | ||
| 198 | + fp8_nd, e8_dn, _, _ = quant_bf16_to_mxfp8_dn(src, m, n_pad) | ||
| 199 | + | ||
| 200 | + group_max = get_group_max_dn(src, group_size=32) | ||
| 201 | + golden_group_max_f32 = group_max.astype(np.float32).view(np.uint32) | ||
| 202 | + | ||
| 203 | + out_dir = os.path.join(GOLDEN_DIR, case_name) | ||
| 204 | + input_bytes = src.astype(np.float32).view(np.uint32).reshape(-1).tobytes() | ||
| 205 | + _write_golden(out_dir, input_bytes, fp8_nd, e8_dn, golden_group_max_f32.reshape(-1).tobytes()) | ||
| 206 | + | ||
| 207 | + | ||
| 208 | +if __name__ == "__main__": | ||
| 209 | + np.random.seed(42) | ||
| 210 | + for case_name, m, n in CASE_PARAMS: | ||
| 211 | + print(f"Generating {case_name}...") | ||
| 212 | + gen_golden_data(case_name, m, n) | ||
| 213 | + for case_name, m, n in FP32_CASE_PARAMS: | ||
| 214 | + print(f"Generating {case_name}...") | ||
| 215 | + gen_golden_data_fp32(case_name, m, n) | ||
| 216 | + print("Done.") | ||
| @@ -0,0 +1,300 @@ | |||
| 1 | +/** | ||
| 2 | +Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | +This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | +CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | +Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | +THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | +INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | +See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | +*/ | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | +using namespace std; | ||
| 16 | +using namespace PtoTestCommon; | ||
| 17 | + | ||
| 18 | +namespace TQuantDNTest { | ||
| 19 | + | ||
| 20 | +template <int Stage, int M, int N, int N_pad> | ||
| 21 | +void LaunchTQuantDN(uint16_t *src, int8_t *fp8_nd, uint8_t *e8_dn, int8_t *fp8_nz, uint8_t *e8_zz, uint16_t *max_dn, | ||
| 22 | + void *stream); | ||
| 23 | + | ||
| 24 | +template <int Stage, int M, int N, int N_pad> | ||
| 25 | +void LaunchTQuantDN_fp32(uint32_t *src, int8_t *fp8_nd, uint8_t *e8_dn, int8_t *fp8_nz, uint8_t *e8_zz, | ||
| 26 | + uint32_t *max_dn, void *stream); | ||
| 27 | + | ||
| 28 | +} // namespace TQuantDNTest | ||
| 29 | + | ||
| 30 | +class TQUANTDNTest : public testing::Test { | ||
| 31 | +protected: | ||
| 32 | + void SetUp() override | ||
| 33 | + {} | ||
| 34 | + void TearDown() override | ||
| 35 | + {} | ||
| 36 | +}; | ||
| 37 | + | ||
| 38 | +std::string GetGoldenDir() | ||
| 39 | +{ | ||
| 40 | + const testing::TestInfo *testInfo = testing::UnitTest::GetInstance()->current_test_info(); | ||
| 41 | + const std::string caseName = testInfo->name(); | ||
| 42 | + const std::string suiteName = testInfo->test_suite_name(); | ||
| 43 | + return "../" + suiteName + "." + caseName; | ||
| 44 | +} | ||
| 45 | + | ||
| 46 | +template <typename T> | ||
| 47 | +void ExpectGoldenMatch(const char *stageName, const char *tensorName, const std::vector<T> &golden, | ||
| 48 | + const std::vector<T> &output) | ||
| 49 | +{ | ||
| 50 | + SCOPED_TRACE(stageName); | ||
| 51 | + ASSERT_EQ(golden.size(), output.size()) << tensorName << " size mismatch"; | ||
| 52 | + EXPECT_TRUE(ResultCmp<T>(golden, output, 0.0f)) << stageName << ": " << tensorName << " mismatch vs golden"; | ||
| 53 | +} | ||
| 54 | + | ||
| 55 | +template <int M, int N, int N_pad> | ||
| 56 | +void test_tquant_dn_bf16() | ||
| 57 | +{ | ||
| 58 | + constexpr int grpSize = 32; | ||
| 59 | + constexpr int hatM = M / grpSize; | ||
| 60 | + constexpr int paddedCols = N_pad; | ||
| 61 | + constexpr int paddedRows16 = ((M + 15) / 16) * 16; | ||
| 62 | + constexpr int virtualRow = paddedRows16 + 1; | ||
| 63 | + constexpr int groupedCols = paddedCols / 32; | ||
| 64 | + constexpr int numGroupsFlat = M * groupedCols; | ||
| 65 | + constexpr int numGroupsFlatAligned = ((numGroupsFlat + 31) / 32) * 32; | ||
| 66 | + size_t srcFileSize = M * paddedCols * sizeof(uint16_t); | ||
| 67 | + size_t fp8NDFileSize = M * paddedCols * sizeof(int8_t); | ||
| 68 | + size_t e8DNFileSize = hatM * paddedCols * sizeof(uint8_t); | ||
| 69 | + size_t fp8NZFileSize = virtualRow * paddedCols * sizeof(int8_t); | ||
| 70 | + size_t e8ZZFileSize = numGroupsFlatAligned * sizeof(uint8_t); | ||
| 71 | + size_t maxDNFileSize = hatM * paddedCols * sizeof(uint16_t); | ||
| 72 | + | ||
| 73 | + aclInit(nullptr); | ||
| 74 | + aclrtSetDevice(0); | ||
| 75 | + aclrtStream stream; | ||
| 76 | + aclrtCreateStream(&stream); | ||
| 77 | + | ||
| 78 | + uint8_t *srcHost; | ||
| 79 | + uint8_t *fp8NDHost; | ||
| 80 | + uint8_t *e8DNHost; | ||
| 81 | + uint8_t *fp8NZHost; | ||
| 82 | + uint8_t *e8ZZHost; | ||
| 83 | + uint16_t *maxDNHost; | ||
| 84 | + uint16_t *srcDevice; | ||
| 85 | + int8_t *fp8NDDevice; | ||
| 86 | + uint8_t *e8DNDevice; | ||
| 87 | + int8_t *fp8NZDevice; | ||
| 88 | + uint8_t *e8ZZDevice; | ||
| 89 | + uint16_t *maxDNDevice; | ||
| 90 | + | ||
| 91 | + aclrtMallocHost((void **)(&srcHost), srcFileSize); | ||
| 92 | + aclrtMallocHost((void **)(&fp8NDHost), fp8NDFileSize); | ||
| 93 | + aclrtMallocHost((void **)(&e8DNHost), e8DNFileSize); | ||
| 94 | + aclrtMallocHost((void **)(&fp8NZHost), fp8NZFileSize); | ||
| 95 | + aclrtMallocHost((void **)(&e8ZZHost), e8ZZFileSize); | ||
| 96 | + aclrtMallocHost((void **)(&maxDNHost), maxDNFileSize); | ||
| 97 | + | ||
| 98 | + aclrtMalloc((void **)&srcDevice, srcFileSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 99 | + aclrtMalloc((void **)&fp8NDDevice, fp8NDFileSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 100 | + aclrtMalloc((void **)&e8DNDevice, e8DNFileSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 101 | + aclrtMalloc((void **)&fp8NZDevice, fp8NZFileSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 102 | + aclrtMalloc((void **)&e8ZZDevice, e8ZZFileSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 103 | + aclrtMalloc((void **)&maxDNDevice, maxDNFileSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 104 | + | ||
| 105 | + ReadFile(GetGoldenDir() + "/input.bin", srcFileSize, srcHost, srcFileSize); | ||
| 106 | + aclrtMemcpy(srcDevice, srcFileSize, srcHost, srcFileSize, ACL_MEMCPY_HOST_TO_DEVICE); | ||
| 107 | + | ||
| 108 | + const std::string goldenDir = GetGoldenDir(); | ||
| 109 | + | ||
| 110 | + // Stage 1: after TQUANT — FP8 ND + E8M0 DN + per-group max | ||
| 111 | + TQuantDNTest::LaunchTQuantDN<1, M, N, N_pad>(srcDevice, fp8NDDevice, e8DNDevice, nullptr, nullptr, maxDNDevice, | ||
| 112 | + stream); | ||
| 113 | + aclError syncRet = aclrtSynchronizeStream(stream); | ||
| 114 | + ASSERT_EQ(syncRet, ACL_SUCCESS) << "Stage1 sync failed: " << aclGetRecentErrMsg(); | ||
| 115 | + | ||
| 116 | + aclrtMemcpy(fp8NDHost, fp8NDFileSize, fp8NDDevice, fp8NDFileSize, ACL_MEMCPY_DEVICE_TO_HOST); | ||
| 117 | + aclrtMemcpy(e8DNHost, e8DNFileSize, e8DNDevice, e8DNFileSize, ACL_MEMCPY_DEVICE_TO_HOST); | ||
| 118 | + aclrtMemcpy(maxDNHost, maxDNFileSize, maxDNDevice, maxDNFileSize, ACL_MEMCPY_DEVICE_TO_HOST); | ||
| 119 | + WriteFile(goldenDir + "/output_fp8_nd.bin", fp8NDHost, fp8NDFileSize); | ||
| 120 | + WriteFile(goldenDir + "/output_e8_dn.bin", e8DNHost, e8DNFileSize); | ||
| 121 | + WriteFile(goldenDir + "/output_group_max.bin", maxDNHost, maxDNFileSize); | ||
| 122 | + | ||
| 123 | + std::vector<uint8_t> goldenFp8Nd(fp8NDFileSize); | ||
| 124 | + std::vector<uint8_t> goldenE8Dn(e8DNFileSize); | ||
| 125 | + std::vector<uint16_t> goldenGroupMax(maxDNFileSize / sizeof(uint16_t)); | ||
| 126 | + std::vector<uint8_t> outFp8Nd(fp8NDFileSize); | ||
| 127 | + std::vector<uint8_t> outE8Dn(e8DNFileSize); | ||
| 128 | + std::vector<uint16_t> outGroupMax(maxDNFileSize / sizeof(uint16_t)); | ||
| 129 | + ReadFile(goldenDir + "/golden_fp8_nd.bin", fp8NDFileSize, goldenFp8Nd.data(), fp8NDFileSize); | ||
| 130 | + ReadFile(goldenDir + "/golden_e8_dn.bin", e8DNFileSize, goldenE8Dn.data(), e8DNFileSize); | ||
| 131 | + ReadFile(goldenDir + "/golden_group_max.bin", maxDNFileSize, goldenGroupMax.data(), maxDNFileSize); | ||
| 132 | + ReadFile(goldenDir + "/output_fp8_nd.bin", fp8NDFileSize, outFp8Nd.data(), fp8NDFileSize); | ||
| 133 | + ReadFile(goldenDir + "/output_e8_dn.bin", e8DNFileSize, outE8Dn.data(), e8DNFileSize); | ||
| 134 | + ReadFile(goldenDir + "/output_group_max.bin", maxDNFileSize, outGroupMax.data(), maxDNFileSize); | ||
| 135 | + ExpectGoldenMatch("Stage1_AfterTQuant", "fp8_nd", goldenFp8Nd, outFp8Nd); | ||
| 136 | + ExpectGoldenMatch("Stage1_AfterTQuant", "e8_dn (exponents)", goldenE8Dn, outE8Dn); | ||
| 137 | + ExpectGoldenMatch("Stage1_AfterTQuant", "group_max", goldenGroupMax, outGroupMax); | ||
| 138 | + | ||
| 139 | + aclrtFree(srcDevice); | ||
| 140 | + aclrtFree(fp8NDDevice); | ||
| 141 | + aclrtFree(e8DNDevice); | ||
| 142 | + aclrtFree(fp8NZDevice); | ||
| 143 | + aclrtFree(e8ZZDevice); | ||
| 144 | + aclrtFree(maxDNDevice); | ||
| 145 | + aclrtFreeHost(srcHost); | ||
| 146 | + aclrtFreeHost(fp8NDHost); | ||
| 147 | + aclrtFreeHost(e8DNHost); | ||
| 148 | + aclrtFreeHost(fp8NZHost); | ||
| 149 | + aclrtFreeHost(e8ZZHost); | ||
| 150 | + aclrtFreeHost(maxDNHost); | ||
| 151 | + aclrtDestroyStream(stream); | ||
| 152 | + aclrtResetDevice(0); | ||
| 153 | + aclFinalize(); | ||
| 154 | +} | ||
| 155 | + | ||
| 156 | +TEST_F(TQUANTDNTest, case_bf16_64x64) | ||
| 157 | +{ | ||
| 158 | + test_tquant_dn_bf16<64, 64, 64>(); | ||
| 159 | +} | ||
| 160 | +TEST_F(TQUANTDNTest, case_bf16_128x64) | ||
| 161 | +{ | ||
| 162 | + test_tquant_dn_bf16<128, 64, 64>(); | ||
| 163 | +} | ||
| 164 | +TEST_F(TQUANTDNTest, case_bf16_64x128) | ||
| 165 | +{ | ||
| 166 | + test_tquant_dn_bf16<64, 128, 128>(); | ||
| 167 | +} | ||
| 168 | +TEST_F(TQUANTDNTest, case_bf16_128x128) | ||
| 169 | +{ | ||
| 170 | + test_tquant_dn_bf16<128, 128, 128>(); | ||
| 171 | +} | ||
| 172 | +TEST_F(TQUANTDNTest, case_bf16_64x256) | ||
| 173 | +{ | ||
| 174 | + test_tquant_dn_bf16<64, 256, 256>(); | ||
| 175 | +} | ||
| 176 | +TEST_F(TQUANTDNTest, case_bf16_128x256) | ||
| 177 | +{ | ||
| 178 | + test_tquant_dn_bf16<128, 256, 256>(); | ||
| 179 | +} | ||
| 180 | +TEST_F(TQUANTDNTest, case_bf16_256x64) | ||
| 181 | +{ | ||
| 182 | + test_tquant_dn_bf16<256, 64, 64>(); | ||
| 183 | +} | ||
| 184 | +TEST_F(TQUANTDNTest, case_bf16_256x128) | ||
| 185 | +{ | ||
| 186 | + test_tquant_dn_bf16<256, 128, 128>(); | ||
| 187 | +} | ||
| 188 | + | ||
| 189 | +template <int M, int N, int N_pad> | ||
| 190 | +void test_tquant_dn_fp32() | ||
| 191 | +{ | ||
| 192 | + constexpr int grpSize = 32; | ||
| 193 | + constexpr int hatM = M / grpSize; | ||
| 194 | + constexpr int paddedCols = N_pad; | ||
| 195 | + constexpr int paddedRows16 = ((M + 15) / 16) * 16; | ||
| 196 | + constexpr int virtualRow = paddedRows16 + 1; | ||
| 197 | + constexpr int groupedCols = paddedCols / 32; | ||
| 198 | + constexpr int numGroupsFlat = M * groupedCols; | ||
| 199 | + constexpr int numGroupsFlatAligned = ((numGroupsFlat + 31) / 32) * 32; | ||
| 200 | + size_t srcFileSize = M * paddedCols * sizeof(uint32_t); | ||
| 201 | + size_t fp8NDFileSize = M * paddedCols * sizeof(int8_t); | ||
| 202 | + size_t e8DNFileSize = hatM * paddedCols * sizeof(uint8_t); | ||
| 203 | + size_t fp8NZFileSize = virtualRow * paddedCols * sizeof(int8_t); | ||
| 204 | + size_t e8ZZFileSize = numGroupsFlatAligned * sizeof(uint8_t); | ||
| 205 | + size_t maxDNFileSize = hatM * paddedCols * sizeof(uint32_t); | ||
| 206 | + | ||
| 207 | + aclInit(nullptr); | ||
| 208 | + aclrtSetDevice(0); | ||
| 209 | + aclrtStream stream; | ||
| 210 | + aclrtCreateStream(&stream); | ||
| 211 | + | ||
| 212 | + uint8_t *srcHost; | ||
| 213 | + uint8_t *fp8NDHost; | ||
| 214 | + uint8_t *e8DNHost; | ||
| 215 | + uint8_t *fp8NZHost; | ||
| 216 | + uint8_t *e8ZZHost; | ||
| 217 | + uint32_t *maxDNHost; | ||
| 218 | + uint32_t *srcDevice; | ||
| 219 | + int8_t *fp8NDDevice; | ||
| 220 | + uint8_t *e8DNDevice; | ||
| 221 | + int8_t *fp8NZDevice; | ||
| 222 | + uint8_t *e8ZZDevice; | ||
| 223 | + uint32_t *maxDNDevice; | ||
| 224 | + | ||
| 225 | + aclrtMallocHost((void **)(&srcHost), srcFileSize); | ||
| 226 | + aclrtMallocHost((void **)(&fp8NDHost), fp8NDFileSize); | ||
| 227 | + aclrtMallocHost((void **)(&e8DNHost), e8DNFileSize); | ||
| 228 | + aclrtMallocHost((void **)(&fp8NZHost), fp8NZFileSize); | ||
| 229 | + aclrtMallocHost((void **)(&e8ZZHost), e8ZZFileSize); | ||
| 230 | + aclrtMallocHost((void **)(&maxDNHost), maxDNFileSize); | ||
| 231 | + | ||
| 232 | + aclrtMalloc((void **)&srcDevice, srcFileSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 233 | + aclrtMalloc((void **)&fp8NDDevice, fp8NDFileSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 234 | + aclrtMalloc((void **)&e8DNDevice, e8DNFileSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 235 | + aclrtMalloc((void **)&fp8NZDevice, fp8NZFileSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 236 | + aclrtMalloc((void **)&e8ZZDevice, e8ZZFileSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 237 | + aclrtMalloc((void **)&maxDNDevice, maxDNFileSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 238 | + | ||
| 239 | + ReadFile(GetGoldenDir() + "/input.bin", srcFileSize, srcHost, srcFileSize); | ||
| 240 | + aclrtMemcpy(srcDevice, srcFileSize, srcHost, srcFileSize, ACL_MEMCPY_HOST_TO_DEVICE); | ||
| 241 | + | ||
| 242 | + const std::string goldenDir = GetGoldenDir(); | ||
| 243 | + | ||
| 244 | + TQuantDNTest::LaunchTQuantDN_fp32<1, M, N, N_pad>(srcDevice, fp8NDDevice, e8DNDevice, nullptr, nullptr, maxDNDevice, | ||
| 245 | + stream); | ||
| 246 | + aclError syncRet = aclrtSynchronizeStream(stream); | ||
| 247 | + ASSERT_EQ(syncRet, ACL_SUCCESS) << "Stage1 sync failed: " << aclGetRecentErrMsg(); | ||
| 248 | + | ||
| 249 | + aclrtMemcpy(fp8NDHost, fp8NDFileSize, fp8NDDevice, fp8NDFileSize, ACL_MEMCPY_DEVICE_TO_HOST); | ||
| 250 | + aclrtMemcpy(e8DNHost, e8DNFileSize, e8DNDevice, e8DNFileSize, ACL_MEMCPY_DEVICE_TO_HOST); | ||
| 251 | + aclrtMemcpy(maxDNHost, maxDNFileSize, maxDNDevice, maxDNFileSize, ACL_MEMCPY_DEVICE_TO_HOST); | ||
| 252 | + WriteFile(goldenDir + "/output_fp8_nd.bin", fp8NDHost, fp8NDFileSize); | ||
| 253 | + WriteFile(goldenDir + "/output_e8_dn.bin", e8DNHost, e8DNFileSize); | ||
| 254 | + WriteFile(goldenDir + "/output_group_max.bin", maxDNHost, maxDNFileSize); | ||
| 255 | + | ||
| 256 | + std::vector<uint8_t> goldenFp8Nd(fp8NDFileSize); | ||
| 257 | + std::vector<uint8_t> goldenE8Dn(e8DNFileSize); | ||
| 258 | + std::vector<uint32_t> goldenGroupMax(maxDNFileSize / sizeof(uint32_t)); | ||
| 259 | + std::vector<uint8_t> outFp8Nd(fp8NDFileSize); | ||
| 260 | + std::vector<uint8_t> outE8Dn(e8DNFileSize); | ||
| 261 | + std::vector<uint32_t> outGroupMax(maxDNFileSize / sizeof(uint32_t)); | ||
| 262 | + ReadFile(goldenDir + "/golden_fp8_nd.bin", fp8NDFileSize, goldenFp8Nd.data(), fp8NDFileSize); | ||
| 263 | + ReadFile(goldenDir + "/golden_e8_dn.bin", e8DNFileSize, goldenE8Dn.data(), e8DNFileSize); | ||
| 264 | + ReadFile(goldenDir + "/golden_group_max.bin", maxDNFileSize, goldenGroupMax.data(), maxDNFileSize); | ||
| 265 | + ReadFile(goldenDir + "/output_fp8_nd.bin", fp8NDFileSize, outFp8Nd.data(), fp8NDFileSize); | ||
| 266 | + ReadFile(goldenDir + "/output_e8_dn.bin", e8DNFileSize, outE8Dn.data(), e8DNFileSize); | ||
| 267 | + ReadFile(goldenDir + "/output_group_max.bin", maxDNFileSize, outGroupMax.data(), maxDNFileSize); | ||
| 268 | + ExpectGoldenMatch("Stage1_AfterTQuant", "fp8_nd", goldenFp8Nd, outFp8Nd); | ||
| 269 | + ExpectGoldenMatch("Stage1_AfterTQuant", "e8_dn (exponents)", goldenE8Dn, outE8Dn); | ||
| 270 | + ExpectGoldenMatch("Stage1_AfterTQuant", "group_max", goldenGroupMax, outGroupMax); | ||
| 271 | + | ||
| 272 | + aclrtFree(srcDevice); | ||
| 273 | + aclrtFree(fp8NDDevice); | ||
| 274 | + aclrtFree(e8DNDevice); | ||
| 275 | + aclrtFree(fp8NZDevice); | ||
| 276 | + aclrtFree(e8ZZDevice); | ||
| 277 | + aclrtFree(maxDNDevice); | ||
| 278 | + aclrtFreeHost(srcHost); | ||
| 279 | + aclrtFreeHost(fp8NDHost); | ||
| 280 | + aclrtFreeHost(e8DNHost); | ||
| 281 | + aclrtFreeHost(fp8NZHost); | ||
| 282 | + aclrtFreeHost(e8ZZHost); | ||
| 283 | + aclrtFreeHost(maxDNHost); | ||
| 284 | + aclrtDestroyStream(stream); | ||
| 285 | + aclrtResetDevice(0); | ||
| 286 | + aclFinalize(); | ||
| 287 | +} | ||
| 288 | + | ||
| 289 | +TEST_F(TQUANTDNTest, case_fp32_64x128) | ||
| 290 | +{ | ||
| 291 | + test_tquant_dn_fp32<64, 128, 128>(); | ||
| 292 | +} | ||
| 293 | +TEST_F(TQUANTDNTest, case_fp32_128x128) | ||
| 294 | +{ | ||
| 295 | + test_tquant_dn_fp32<128, 128, 128>(); | ||
| 296 | +} | ||
| 297 | +TEST_F(TQUANTDNTest, case_fp32_64x256) | ||
| 298 | +{ | ||
| 299 | + test_tquant_dn_fp32<64, 256, 256>(); | ||
| 300 | +} | ||
| @@ -0,0 +1,223 @@ | |||
| 1 | +/** | ||
| 2 | +Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | +This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | +CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | +Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | +THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | +INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | +See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | +*/ | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | +using namespace pto; | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + | ||
| 20 | +namespace TQuantDNTest { | ||
| 21 | + | ||
| 22 | +// Stage 1: after TQUANT (FP8 ND + E8M0 DN) | ||
| 23 | +// Stage 2: after TQUANT + FP8 ND->NZ | ||
| 24 | +// Stage 3: full pipeline including E8 DN->ZZ | ||
| 25 | +template <int Stage, typename T, int M, int N, int N_pad> | ||
| 26 | +__global__ AICORE void runTQuantDN(__gm__ T __in__ *src_gm, __gm__ int8_t __out__ *fp8_nd_gm, | ||
| 27 | + __gm__ uint8_t __out__ *e8_dn_gm, __gm__ int8_t __out__ *fp8_nz_gm, | ||
| 28 | + __gm__ uint8_t __out__ *e8_zz_gm, __gm__ T __out__ *max_dn_gm) | ||
| 29 | +{ | ||
| 30 | + static_assert(Stage >= 1 && Stage <= 3, "Stage must be 1 (quant), 2 (nz), or 3 (zz)."); | ||
| 31 | + | ||
| 32 | + constexpr uint32_t grpSize = 32; | ||
| 33 | + constexpr uint32_t hatM = M / grpSize; | ||
| 34 | + constexpr uint32_t paddedCols = N_pad; | ||
| 35 | + constexpr uint32_t groupedColsValid = paddedCols / 32; | ||
| 36 | + constexpr uint32_t numGroupsFlat = M * groupedColsValid; | ||
| 37 | + constexpr uint32_t numGroupsFlatAligned = PTO_CEIL(numGroupsFlat, 32); | ||
| 38 | + constexpr uint32_t paddedRows16 = PTO_CEIL(M, FRACTAL_NZ_ROW); | ||
| 39 | + constexpr uint32_t virtualRow = paddedRows16 + 1; | ||
| 40 | + | ||
| 41 | + using SrcTile = | ||
| 42 | + Tile<TileType::Vec, T, M, paddedCols, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Zero>; | ||
| 43 | + using DstFP8Tile = Tile<TileType::Vec, int8_t, M, paddedCols, BLayout::RowMajor, M, paddedCols, SLayout::NoneBox, | ||
| 44 | + 512, PadValue::Zero>; | ||
| 45 | + using MaxTile = | ||
| 46 | + Tile<TileType::Vec, T, hatM, paddedCols, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Zero>; | ||
| 47 | + using ScalingTile = | ||
| 48 | + Tile<TileType::Vec, T, hatM, paddedCols, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Zero>; | ||
| 49 | + using E8NdTile = Tile<TileType::Vec, uint8_t, hatM, paddedCols, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, | ||
| 50 | + PadValue::Zero>; | ||
| 51 | + | ||
| 52 | + using E8DnTile = Tile<TileType::Vec, uint8_t, hatM, paddedCols, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, | ||
| 53 | + PadValue::Zero>; | ||
| 54 | + | ||
| 55 | + using E8ZzTile = Tile<TileType::Vec, uint8_t, paddedRows16, groupedColsValid, BLayout::RowMajor, -1, -1, | ||
| 56 | + SLayout::RowMajor, 32, PadValue::Zero>; | ||
| 57 | + using E8StoreTile = Tile<TileType::Vec, uint8_t, 1, numGroupsFlatAligned, BLayout::RowMajor, -1, -1, | ||
| 58 | + SLayout::NoneBox, 512, PadValue::Zero>; | ||
| 59 | + | ||
| 60 | + using Fp8NZTile = Tile<TileType::Vec, int8_t, virtualRow, paddedCols, BLayout::ColMajor, M, paddedCols, | ||
| 61 | + SLayout::RowMajor, 512, PadValue::Null, CompactMode::RowPlusOne>; | ||
| 62 | + | ||
| 63 | + constexpr uint32_t colBlkCount = paddedCols / 16; | ||
| 64 | + constexpr uint32_t hatP = hatM / 2; | ||
| 65 | + constexpr uint32_t tmpBufSize = | ||
| 66 | + (BLOCK_SIZE / sizeof(uint16_t) + | ||
| 67 | + (colBlkCount > hatP ? colBlkCount : hatP) * (hatP > colBlkCount ? hatP : colBlkCount) + | ||
| 68 | + BLOCK_SIZE / sizeof(uint16_t)) * | ||
| 69 | + sizeof(uint16_t); | ||
| 70 | + constexpr uint32_t tmpBufSizeAligned = PTO_CEIL(tmpBufSize, 32); | ||
| 71 | + | ||
| 72 | + using TmpTile = Tile<TileType::Vec, uint8_t, 1, tmpBufSizeAligned, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, | ||
| 73 | + PadValue::Zero>; | ||
| 74 | + | ||
| 75 | + SrcTile srcTile(M, paddedCols); | ||
| 76 | + DstFP8Tile fp8Tile; | ||
| 77 | + MaxTile maxPerGpTile(hatM, paddedCols); | ||
| 78 | + ScalingTile scalingTile(hatM, paddedCols); | ||
| 79 | + E8NdTile e8Tile(hatM, paddedCols); | ||
| 80 | + E8DnTile e8DnTile(hatM, paddedCols); | ||
| 81 | + E8ZzTile e8ZzTile(paddedRows16, groupedColsValid); | ||
| 82 | + E8StoreTile e8StoreTile(1, numGroupsFlatAligned); | ||
| 83 | + Fp8NZTile fp8TileNZ; | ||
| 84 | + TmpTile tmpTile(1, tmpBufSizeAligned); | ||
| 85 | + | ||
| 86 | + using SrcGlobal = GlobalTensor<T, Shape<1, 1, 1, M, N_pad>, pto::Stride<1, 1, 1, N_pad, 1>>; | ||
| 87 | + SrcGlobal srcGlobal(src_gm); | ||
| 88 | + | ||
| 89 | + using DstFp8NdGlobal = GlobalTensor<int8_t, Shape<1, 1, 1, M, paddedCols>, pto::Stride<1, 1, 1, paddedCols, 1>>; | ||
| 90 | + DstFp8NdGlobal fp8NdGlobal(fp8_nd_gm); | ||
| 91 | + | ||
| 92 | + using DstMaxGlobal = GlobalTensor<T, Shape<1, 1, 1, hatM, paddedCols>, pto::Stride<1, 1, 1, paddedCols, 1>>; | ||
| 93 | + DstMaxGlobal maxGlobal(max_dn_gm); | ||
| 94 | + | ||
| 95 | + using DstE8DnGlobal = GlobalTensor<uint8_t, Shape<1, 1, 1, hatM, paddedCols>, pto::Stride<1, 1, 1, paddedCols, 1>>; | ||
| 96 | + DstE8DnGlobal e8DnGlobal(e8_dn_gm); | ||
| 97 | + | ||
| 98 | + using DstE8Global = | ||
| 99 | + GlobalTensor<uint8_t, Shape<1, 1, 1, 1, numGroupsFlatAligned>, pto::Stride<1, 1, 1, numGroupsFlatAligned, 1>>; | ||
| 100 | + DstE8Global e8Global(e8_zz_gm); | ||
| 101 | + | ||
| 102 | + using DstFp8GlobalNZ = GlobalTensor<int8_t, TileShape2D<int8_t, M, paddedCols, Layout::NZ>, | ||
| 103 | + BaseShape2D<int8_t, M, paddedCols, Layout::NZ>, Layout::NZ>; | ||
| 104 | + DstFp8GlobalNZ fp8GlobalNZ((__gm__ int8_t *)fp8_nz_gm); | ||
| 105 | + | ||
| 106 | + constexpr uint32_t srcTileBytes = M * paddedCols * sizeof(T); | ||
| 107 | + constexpr uint32_t maxTileBytes = hatM * paddedCols * sizeof(T); | ||
| 108 | + constexpr uint32_t scalingTileBytes = hatM * paddedCols * sizeof(T); | ||
| 109 | + constexpr uint32_t e8TileBytes = hatM * paddedCols; | ||
| 110 | + constexpr uint32_t e8DnTileBytes = hatM * paddedCols; | ||
| 111 | + constexpr uint32_t fp8TileBytes = M * paddedCols; | ||
| 112 | + | ||
| 113 | + // Keep source and destination UB tiles orthogonal: place fp8Tile after all | ||
| 114 | + // input/work tiles so TQuant reads src and writes dst to non-overlapping | ||
| 115 | + // regions (on-board store ordering is not guaranteed). | ||
| 116 | + constexpr uint32_t srcTileAddr = 0x0; | ||
| 117 | + constexpr uint32_t maxTileAddr = PTO_CEIL(srcTileAddr + srcTileBytes, 0x20); | ||
| 118 | + constexpr uint32_t scalingTileAddr = PTO_CEIL(maxTileAddr + maxTileBytes, 0x20); | ||
| 119 | + constexpr uint32_t e8TileAddr = PTO_CEIL(scalingTileAddr + scalingTileBytes, 0x20); | ||
| 120 | + constexpr uint32_t e8DnTileAddr = PTO_CEIL(e8TileAddr + e8TileBytes, 0x20); | ||
| 121 | + constexpr uint32_t fp8TileAddr = PTO_CEIL(e8DnTileAddr + e8DnTileBytes, 0x20); | ||
| 122 | + constexpr uint32_t C0_SIZE_B = 32; | ||
| 123 | + constexpr uint32_t nColGroupsNZ = paddedCols / C0_SIZE_B; | ||
| 124 | + constexpr uint32_t fp8NZTileBytes = | ||
| 125 | + (nColGroupsNZ > 1) ? (nColGroupsNZ - 1) * (paddedRows16 + 1) * C0_SIZE_B + paddedRows16 * C0_SIZE_B : | ||
| 126 | + paddedRows16 * C0_SIZE_B; | ||
| 127 | + constexpr uint32_t fp8NZTileAddr = PTO_CEIL(fp8TileAddr + fp8TileBytes, 0x20); | ||
| 128 | + // workTileEnd marks the end of the TQuant input-side tiles; fp8Tile now lives | ||
| 129 | + // after it, so use fp8NZEnd to find where the ZZ/tmp scratch area can start. | ||
| 130 | + constexpr uint32_t workTileEnd = e8DnTileAddr + e8DnTileBytes; | ||
| 131 | + constexpr uint32_t fp8NZEnd = fp8NZTileAddr + fp8NZTileBytes; | ||
| 132 | + constexpr uint32_t zzTmpStart = PTO_CEIL(workTileEnd > fp8NZEnd ? workTileEnd : fp8NZEnd, 0x20); | ||
| 133 | + constexpr uint32_t e8ZzTileAddr = zzTmpStart; | ||
| 134 | + constexpr uint32_t e8StoreTileAddr = zzTmpStart; | ||
| 135 | + constexpr uint32_t tmpTileAddr = PTO_CEIL(e8ZzTileAddr + numGroupsFlatAligned, 0x20); | ||
| 136 | + constexpr uint32_t layoutEnd = PTO_CEIL(tmpTileAddr + tmpBufSizeAligned, 0x100); | ||
| 137 | + static_assert(layoutEnd <= 0x40000, "UB layout exceeds 256 KB."); | ||
| 138 | + | ||
| 139 | + TASSIGN(srcTile, srcTileAddr); | ||
| 140 | + TASSIGN(maxPerGpTile, maxTileAddr); | ||
| 141 | + TASSIGN(scalingTile, scalingTileAddr); | ||
| 142 | + TASSIGN(e8Tile, e8TileAddr); | ||
| 143 | + TASSIGN(e8DnTile, e8DnTileAddr); | ||
| 144 | + TASSIGN(e8ZzTile, e8ZzTileAddr); | ||
| 145 | + TASSIGN(e8StoreTile, e8StoreTileAddr); | ||
| 146 | + TASSIGN(fp8Tile, fp8TileAddr); | ||
| 147 | + TASSIGN(fp8TileNZ, fp8NZTileAddr); | ||
| 148 | + TASSIGN(tmpTile, tmpTileAddr); | ||
| 149 | + | ||
| 150 | + TLOAD(srcTile, srcGlobal); | ||
| 151 | + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); | ||
| 152 | + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); | ||
| 153 | + | ||
| 154 | + TQUANT<QuantType::MXFP8, QuantScaleAlg::OCP>(fp8Tile, srcTile, &e8Tile, &maxPerGpTile, &scalingTile, &e8DnTile); | ||
| 155 | + | ||
| 156 | + // TQuant writes the exponent tile (e8Tile) in row-major [hatM, paddedCols]. | ||
| 157 | + // Copy it to e8DnTile so the UB tile shape matches the GM shape exactly. | ||
| 158 | + TMOV(e8DnTile, e8Tile); | ||
| 159 | + | ||
| 160 | + if constexpr (Stage == 1) { | ||
| 161 | + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); | ||
| 162 | + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); | ||
| 163 | + TSTORE(fp8NdGlobal, fp8Tile); | ||
| 164 | + TSTORE(e8DnGlobal, e8DnTile); | ||
| 165 | + TSTORE(maxGlobal, maxPerGpTile); | ||
| 166 | + return; | ||
| 167 | + } | ||
| 168 | + | ||
| 169 | + TMOV(fp8TileNZ, fp8Tile); | ||
| 170 | + | ||
| 171 | + if constexpr (Stage == 2) { | ||
| 172 | + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); | ||
| 173 | + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); | ||
| 174 | + TSTORE(fp8GlobalNZ, fp8TileNZ); | ||
| 175 | + return; | ||
| 176 | + } | ||
| 177 | + | ||
| 178 | + TMOV(e8ZzTile, e8DnTile, tmpTile); | ||
| 179 | + | ||
| 180 | + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); | ||
| 181 | + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); | ||
| 182 | + TSTORE(e8Global, e8StoreTile); | ||
| 183 | + TSTORE(fp8GlobalNZ, fp8TileNZ); | ||
| 184 | +} | ||
| 185 | + | ||
| 186 | +template <int Stage, int M, int N, int N_pad> | ||
| 187 | +void LaunchTQuantDN(uint16_t *src, int8_t *fp8_nd, uint8_t *e8_dn, int8_t *fp8_nz, uint8_t *e8_zz, uint16_t *max_dn, | ||
| 188 | + void *stream) | ||
| 189 | +{ | ||
| 190 | + runTQuantDN<Stage, bfloat16_t, M, N, N_pad> | ||
| 191 | + <<<1, nullptr, stream>>>((bfloat16_t *)src, fp8_nd, e8_dn, fp8_nz, e8_zz, (bfloat16_t *)max_dn); | ||
| 192 | +} | ||
| 193 | + | ||
| 194 | +template <int Stage, int M, int N, int N_pad> | ||
| 195 | +void LaunchTQuantDN_fp32(uint32_t *src, int8_t *fp8_nd, uint8_t *e8_dn, int8_t *fp8_nz, uint8_t *e8_zz, | ||
| 196 | + uint32_t *max_dn, void *stream) | ||
| 197 | +{ | ||
| 198 | + runTQuantDN<Stage, float, M, N, N_pad> | ||
| 199 | + <<<1, nullptr, stream>>>((float *)src, fp8_nd, e8_dn, fp8_nz, e8_zz, (float *)max_dn); | ||
| 200 | +} | ||
| 201 | + | ||
| 202 | + | ||
| 203 | + template void LaunchTQuantDN<S, M, N, NP>(uint16_t *, int8_t *, uint8_t *, int8_t *, uint8_t *, uint16_t *, void *) | ||
| 204 | + | ||
| 205 | + | ||
| 206 | + template void LaunchTQuantDN_fp32<S, M, N, NP>(uint32_t *, int8_t *, uint8_t *, int8_t *, uint8_t *, uint32_t *, \ | ||
| 207 | + void *) | ||
| 208 | + | ||
| 209 | +INSTANTIATE_TQUANT_DN_STAGE(1, 128, 128, 128); | ||
| 210 | +INSTANTIATE_TQUANT_DN_STAGE(1, 64, 128, 128); | ||
| 211 | +INSTANTIATE_TQUANT_DN_STAGE(1, 64, 256, 256); | ||
| 212 | +INSTANTIATE_TQUANT_DN_STAGE(1, 128, 256, 256); | ||
| 213 | +INSTANTIATE_TQUANT_DN_STAGE(1, 64, 64, 64); | ||
| 214 | +INSTANTIATE_TQUANT_DN_STAGE(1, 128, 64, 64); | ||
| 215 | +INSTANTIATE_TQUANT_DN_STAGE(1, 256, 64, 64); | ||
| 216 | +INSTANTIATE_TQUANT_DN_STAGE(1, 256, 128, 128); | ||
| 217 | +INSTANTIATE_TQUANT_DN_STAGE_FP32(1, 64, 128, 128); | ||
| 218 | +INSTANTIATE_TQUANT_DN_STAGE_FP32(1, 128, 128, 128); | ||
| 219 | +INSTANTIATE_TQUANT_DN_STAGE_FP32(1, 64, 256, 256); | ||
| 220 | + | ||
| 221 | + | ||
| 222 | + | ||
| 223 | +} // namespace TQuantDNTest | ||
| @@ -850,18 +850,18 @@ if [ "$ENABLE_KIRIN9030" = "true" ]; then | |||
| 850 | python3 tests/script/run_st.py $ARGS -w -v kirin9030 -t texpands_mat | 850 | python3 tests/script/run_st.py $ARGS -w -v kirin9030 -t texpands_mat |
| 851 | fi | 851 | fi |
| 852 | 852 | ||
| 853 | -if [ "$ENABLE_COMM" = "true" ]; then | 853 | +# if [ "$ENABLE_COMM" = "true" ]; then |
| 854 | - if [ "$ENABLE_A3" = "true" ]; then | 854 | +# if [ "$ENABLE_A3" = "true" ]; then |
| 855 | - python3 tests/script/run_st.py $ARGS -v a3 -t comm/tput | 855 | +# python3 tests/script/run_st.py $ARGS -v a3 -t comm/tput |
| 856 | - python3 tests/script/run_st.py $ARGS -v a3 -t comm/tget | 856 | +# python3 tests/script/run_st.py $ARGS -v a3 -t comm/tget |
| 857 | - python3 tests/script/run_st.py $ARGS -v a3 -t comm/tnotify | 857 | +# python3 tests/script/run_st.py $ARGS -v a3 -t comm/tnotify |
| 858 | - python3 tests/script/run_st.py $ARGS -v a3 -t comm/twait | 858 | +# python3 tests/script/run_st.py $ARGS -v a3 -t comm/twait |
| 859 | - python3 tests/script/run_st.py $ARGS -v a3 -t comm/ttest | 859 | +# python3 tests/script/run_st.py $ARGS -v a3 -t comm/ttest |
| 860 | - python3 tests/script/run_st.py $ARGS -v a3 -t comm/tgather | 860 | +# python3 tests/script/run_st.py $ARGS -v a3 -t comm/tgather |
| 861 | - python3 tests/script/run_st.py $ARGS -v a3 -t comm/tscatter | 861 | +# python3 tests/script/run_st.py $ARGS -v a3 -t comm/tscatter |
| 862 | - python3 tests/script/run_st.py $ARGS -v a3 -t comm/treduce | 862 | +# python3 tests/script/run_st.py $ARGS -v a3 -t comm/treduce |
| 863 | - python3 tests/script/run_st.py $ARGS -v a3 -t comm/tbroadcast | 863 | +# python3 tests/script/run_st.py $ARGS -v a3 -t comm/tbroadcast |
| 864 | - python3 tests/script/run_st.py $ARGS -v a3 -t comm/tput_async | 864 | +# python3 tests/script/run_st.py $ARGS -v a3 -t comm/tput_async |
| 865 | - python3 tests/script/run_st.py $ARGS -v a3 -t comm/tget_async | 865 | +# python3 tests/script/run_st.py $ARGS -v a3 -t comm/tget_async |
| 866 | - fi | 866 | +# fi |
| 867 | -fi | 867 | +# fi |


🔵 Low Priority
变更位置:第 2392 行
MaskReg preg_b16 = CreatePredicate<bfloat16_t>(preg_cols_b16);函数
calcQuantizedFP8Values_DN_B16是模板,T可以是bfloat16_t或half。此处硬编码CreatePredicate<bfloat16_t>,当T = half时谓词类型与数据不匹配。对比 pre-existing 代码第 1006–1007 行
CalcQuantizedFP8Values_B16_Window使用CreatePredicate<T>(...)正确推导为 T。虽然在此硬件上 bf16 和 fp16 共享同一向量寄存器类型且当前未实例化 fp16 用例,但这是潜在的类型 bug,一旦有人添加 fp16 DN 测试即可能触发问题。修复:将
bfloat16_t改为T。建议:将
CreatePredicate<bfloat16_t>改为CreatePredicate<T>,与其他同模式模板函数的写法对齐。