已合并
mgather gm -> l1 a2a3, a5 & simt variants #1136
Sathi Sarveswara Reddy创建于 6月17日
mgather gm -> l1 a2a3, a5 & simt variants #1136
已合并
共 15 个文件变更+1935-92
| @@ -13,6 +13,8 @@ | |||
| 13 | 13 | ||
| 14 | Out-of-bounds handling is selected through the `GatherOOB` template parameter. `MGATHER` has no atomic or conflict policy: every destination slot has exactly one defined source index, so collisions cannot occur. | 14 | Out-of-bounds handling is selected through the `GatherOOB` template parameter. `MGATHER` has no atomic or conflict policy: every destination slot has exactly one defined source index, so collisions cannot occur. |
| 15 | 15 | ||
| 16 | +The destination may also be an **L1 / cube `TileType::Mat` tile in NZ layout** (with the index supplied as a GM tensor). This GM → L1 path — for both `Coalesce::Row` and `Coalesce::Elem`, on A2/A3 and A5 — is documented in the [GM → L1 Gather](#gm--l1-gather-tiletypemat-destination) section below; the GM → UB behaviour described here is unchanged. | ||
| 17 | + | ||
| 16 | Per-target dispatch summary: | 18 | Per-target dispatch summary: |
| 17 | 19 | ||
| 18 | - **CPU Simulator** — pure C++ reference. The implementation walks `validRow * validCol` and reads `table[idx[i, j]]` (Elem semantics); the CPU sim does not have a separate Row coalesce path. Row iteration uses `pto::cpu::parallel_for_rows`, which by default runs sequentially in row-major order because `PTO_CPU_MAX_THREADS` defaults to `1u`. No destination collision is possible for gather, so the iteration order is observationally irrelevant. | 20 | - **CPU Simulator** — pure C++ reference. The implementation walks `validRow * validCol` and reads `table[idx[i, j]]` (Elem semantics); the CPU sim does not have a separate Row coalesce path. Row iteration uses `pto::cpu::parallel_for_rows`, which by default runs sequentially in row-major order because `PTO_CPU_MAX_THREADS` defaults to `1u`. No destination collision is possible for gather, so the iteration order is observationally irrelevant. |
| @@ -648,6 +650,219 @@ AICORE void example_scalar(__gm__ float* tablePtr, __gm__ int32_t* idxPtr) | |||
| 648 | pto.mgather ins(%mem, %idx : !pto.partition_tensor_view<MxNxdtype>, !pto.tile_buf<...>) outs(%dst : !pto.tile_buf<...>) | 650 | pto.mgather ins(%mem, %idx : !pto.partition_tensor_view<MxNxdtype>, !pto.tile_buf<...>) outs(%dst : !pto.tile_buf<...>) |
| 649 | ``` | 651 | ``` |
| 650 | 652 | ||
| 653 | +## GM → L1 Gather (`TileType::Mat` destination) | ||
| 654 | + | ||
| 655 | +In addition to the GM → UB gather described above, `MGATHER` supports gathering an | ||
| 656 | +indexed selection directly into an **L1 / cube `TileType::Mat` tile in NZ fractal layout**. | ||
| 657 | +This is the form a matmul consumes (the `A`/`NZ` operand produced by `TLOAD(MatTile_NZ, | ||
| 658 | +GlobalTensor_ND)`), so a gathered table can be fed straight into `TEXTRACT` / `TMOV` → | ||
| 659 | +`TMATMUL` without a UB round-trip. Both `Coalesce::Row` and `Coalesce::Elem` are supported, | ||
| 660 | +on **A2/A3 and A5**, for all element dtypes and all four `GatherOOB` policies. | ||
| 661 | + | ||
| 662 | +The GM → L1 path is selected automatically when `TileDst::Loc == TileType::Mat`; the GM → UB | ||
| 663 | +behaviour is unchanged for `TileType::Vec` destinations. | ||
| 664 | + | ||
| 665 | +### Index source — GM | ||
| 666 | + | ||
| 667 | +For GM → L1 the scalar that reads the index and issues the L1 DMA runs on the **cube core** | ||
| 668 | +(`__DAV_CUBE__`). On A5 the AIC and AIV are separate cores and the AIC cannot read AIV's UB, | ||
| 669 | +so a UB index tile is not portable. The GM → L1 variant therefore takes the **index as a GM | ||
| 670 | +`GlobalTensor`** (`int32_t` / `uint32_t`), read by the same core that issues the L1 transfer. | ||
| 671 | +This works identically on A2/A3 and A5. | ||
| 672 | + | ||
| 673 | +### API | ||
| 674 | + | ||
| 675 | +Row mode reuses the 3-argument form; the dispatcher routes to the GM → L1 path because the | ||
| 676 | +destination is a `Mat` tile: | ||
| 677 | + | ||
| 678 | +```cpp | ||
| 679 | +template <Coalesce CMode = Coalesce::Row, GatherOOB Oob = GatherOOB::Undefined, | ||
| 680 | + typename MatTileDst, typename GlobalTable, typename GlobalIdx> | ||
| 681 | +PTO_INST RecordEvent MGATHER(MatTileDst& dst, GlobalTable& table, GlobalIdx& idx); | ||
| 682 | +``` | ||
| 683 | + | ||
| 684 | +Elem mode needs a contiguous **GM scratch** workspace to stage the discrete elements into NZ | ||
| 685 | +layout before the bulk GM → L1 copy, so it takes a fourth operand: | ||
| 686 | + | ||
| 687 | +```cpp | ||
| 688 | +template <Coalesce CMode = Coalesce::Elem, GatherOOB Oob = GatherOOB::Undefined, | ||
| 689 | + typename MatTileDst, typename GlobalTable, typename GlobalIdx, typename GlobalScratch> | ||
| 690 | +PTO_INST RecordEvent MGATHER(MatTileDst& dst, GlobalTable& table, GlobalIdx& idx, | ||
| 691 | + GlobalScratch& scratch); | ||
| 692 | +``` | ||
| 693 | + | ||
| 694 | +`scratch` is a GM `GlobalTensor<T, …>` of at least `TileDst::Rows * TileDst::Cols` elements. | ||
| 695 | + | ||
| 696 | +### Algorithm | ||
| 697 | + | ||
| 698 | +- **Row mode.** For each logical row `r`, the scalar reads `idx[r]` from GM, remaps it per the | ||
| 699 | + `GatherOOB` policy, and issues **one `copy_gm_to_cbuf_multi_nd2nz` per row** straight from | ||
| 700 | + the ND table row `table[safeIdx, :]` into the NZ slot `dstBase + r * kC0` of the L1 tile. The | ||
| 701 | + ND → NZ conversion happens in-flight, exactly like the per-row slice of | ||
| 702 | + `TLoadGm2L1Nd2nz` (A2/A3) / `TLoadCubeND2NZ` (A5). No GM scratch is needed. | ||
| 703 | +- **Elem mode.** The scalar gathers each `table[idx[r, c]]` into the GM `scratch` buffer at the | ||
| 704 | + NZ offset `(c / kC0) * (Rows * kC0) + r * kC0 + (c % kC0)` (the buffer is pre-zeroed so OOB / | ||
| 705 | + `Zero` lanes stay zero). After the scalar fill, the buffer is flushed to DDR (`dcci` per | ||
| 706 | + cache line + `dsb(DSB_DDR)`) so the MTE2 engine observes the scalar writes, then a single | ||
| 707 | + contiguous `copy_gm_to_cbuf` (`pto_copy_gm_to_cbuf_align_v2` on A5) moves the whole NZ buffer | ||
| 708 | + GM → L1. | ||
| 709 | + | ||
| 710 | +`kC0 = C0_SIZE_BYTE / sizeof(T) = 32 / sizeof(T)`. `OOB::Zero` in Row mode pre-zeros the whole | ||
| 711 | +L1 tile once via `pto_create_cbuf_matrix` and skips the DMA for OOB rows. | ||
| 712 | + | ||
| 713 | +### Constraints (`MGatherCheckGm2L1`, A2/A3 and A5) | ||
| 714 | + | ||
| 715 | +- **Dtypes.** A2/A3: `int8/uint8/int16/uint16/int32/uint32/half/bfloat16/float`. A5 additionally | ||
| 716 | + allows `hifloat8/float8_e4m3/float8_e5m2`. Row mode further requires `sizeof(T) <= 4` (the | ||
| 717 | + nd2nz engine handles b8/b16/b32). Elem mode supports every listed dtype (byte-wise staging). | ||
| 718 | +- **Index.** `idx` must be a GM `GlobalTensor` of `int32_t` or `uint32_t`. | ||
| 719 | +- **Table.** `GlobalTable::DType == __gm__ T` and `GlobalTable::layout == Layout::ND`. | ||
| 720 | +- **Scratch (Elem only).** `GlobalScratch::DType == __gm__ T`. | ||
| 721 | +- **Destination.** `TileDst::Loc == TileType::Mat`, NZ form | ||
| 722 | + (`!isRowMajor && SFractal == SLayout::RowMajor && SFractalSize == TileConfig::fractalABSize` | ||
| 723 | + = 512 B), `TileDst::Cols % (C0_SIZE_BYTE / sizeof(T)) == 0`, and | ||
| 724 | + `TileDst::Rows % FRACTAL_NZ_ROW (16) == 0`. | ||
| 725 | +- Shape-coupled checks are gated on `if constexpr (DIM > 0)` / `ValidRow|ValidCol > 0`, so both | ||
| 726 | + static and runtime-dynamic shapes are accepted (`Rows` / `Cols` are always compile-time and | ||
| 727 | + govern NZ addressing). `tableRows = ∏ Shape[0..3]`, `tableRowStride = GetStride(DIM_3)` for | ||
| 728 | + Row; `tableSize = ∏ Shape[0..4]` for Elem. | ||
| 729 | + | ||
| 730 | +### Per-arch realisation | ||
| 731 | + | ||
| 732 | +- **A2/A3 (vec-core, unified AI core).** The scalar address computation, the `nd2nz` / | ||
| 733 | + `copy_gm_to_cbuf`, and the optional `pto_create_cbuf_matrix` pre-zero all run on the same core. | ||
| 734 | + Internal handshakes use specific producer→consumer pairs only — never `pipe_barrier(PIPE_ALL)`. | ||
| 735 | + In Row mode the `OOB::Zero` pre-zero (`pto_create_cbuf_matrix`) and the `nd2nz` gather are both | ||
| 736 | + MTE2 instructions on a single in-order DMA queue, so the WAW on L1 is ordered with no extra | ||
| 737 | + flag; Elem mode uses `dcci`/`dsb` + `S→MTE2` before the bulk `copy_gm_to_cbuf`. The 11-argument | ||
| 738 | + `pto_copy_gm_to_cbuf_multi_nd2nz` form is used for Row. | ||
| 739 | +- **A5 (separate AIC + AIV).** The gather (nd2nz / `copy_gm_to_cbuf_align_v2`) runs on the AIC | ||
| 740 | + cube core; `set_mte2_nz_para` configures the NZ destination strides once before the Row loop. | ||
| 741 | + The L1 tile is a cube-only resource, so reading it back to GM (e.g. for verification or a | ||
| 742 | + vector consumer) goes AIC → UB (`copy_cbuf_to_ubuf`) → AIV → GM with an intra-block | ||
| 743 | + handshake, exactly like `tload_mix`. `copy_cbuf_to_gm` / `copy_ubuf_to_gm` are not available | ||
| 744 | + on the AIC cube target. Elem mode additionally offers an opt-in **SIMT executor** | ||
| 745 | + (`GatherExec::Simt`) that runs the gather on the AIV vector core — see | ||
| 746 | + [A5 only — SIMT executor](#a5-only--simt-executor-for-elem-gm--l1-gatherexecsimt) below. | ||
| 747 | + | ||
| 748 | +### Example — Row gather into an L1 NZ tile | ||
| 749 | + | ||
| 750 | +```cpp | ||
| 751 | +template <typename T, int R, int C, int TableRows> | ||
| 752 | +AICORE void example_gm2l1_row(__gm__ T* tablePtr, __gm__ int32_t* idxPtr) | ||
| 753 | +{ | ||
| 754 | + using TableShape = Shape<1, 1, 1, TableRows, C>; | ||
| 755 | + using TableStride = Stride<1, 1, 1, C, 1>; | ||
| 756 | + using IdxShape = Shape<1, 1, 1, 1, R>; | ||
| 757 | + using IdxStride = Stride<1, 1, 1, R, 1>; | ||
| 758 | + GlobalTensor<T, TableShape, TableStride, Layout::ND> tableGM(tablePtr); | ||
| 759 | + GlobalTensor<int32_t, IdxShape, IdxStride, Layout::ND> idxGM(idxPtr); | ||
| 760 | + | ||
| 761 | + using DstTile = Tile<TileType::Mat, T, R, C, BLayout::ColMajor, R, C, SLayout::RowMajor, 512>; | ||
| 762 | + DstTile dst; TASSIGN(dst, 0x0); | ||
| 763 | + | ||
| 764 | + MGATHER<Coalesce::Row, GatherOOB::Clamp>(dst, tableGM, idxGM); // GM (ND) -> L1 (NZ) | ||
| 765 | +} | ||
| 766 | +``` | ||
| 767 | + | ||
| 768 | +### Example — Elem gather into an L1 NZ tile (with GM scratch) | ||
| 769 | + | ||
| 770 | +```cpp | ||
| 771 | +template <typename T, int R, int C, int TableSize> | ||
| 772 | +AICORE void example_gm2l1_elem(__gm__ T* tablePtr, __gm__ int32_t* idxPtr, __gm__ T* scratchPtr) | ||
| 773 | +{ | ||
| 774 | + using TableShape = Shape<1, 1, 1, 1, TableSize>; | ||
| 775 | + using TableStride = Stride<1, 1, 1, TableSize, 1>; | ||
| 776 | + using IdxShape = Shape<1, 1, 1, R, C>; | ||
| 777 | + using IdxStride = Stride<1, 1, 1, C, 1>; | ||
| 778 | + using ScratchShape = Shape<1, 1, 1, 1, R * C>; | ||
| 779 | + using ScratchStride= Stride<1, 1, 1, R * C, 1>; | ||
| 780 | + GlobalTensor<T, TableShape, TableStride, Layout::ND> tableGM(tablePtr); | ||
| 781 | + GlobalTensor<int32_t, IdxShape, IdxStride, Layout::ND> idxGM(idxPtr); | ||
| 782 | + GlobalTensor<T, ScratchShape, ScratchStride, Layout::ND> scratchGM(scratchPtr); | ||
| 783 | + | ||
| 784 | + using DstTile = Tile<TileType::Mat, T, R, C, BLayout::ColMajor, R, C, SLayout::RowMajor, 512>; | ||
| 785 | + DstTile dst; TASSIGN(dst, 0x0); | ||
| 786 | + | ||
| 787 | + MGATHER<Coalesce::Elem, GatherOOB::Zero>(dst, tableGM, idxGM, scratchGM); // GM -> GM scratch (NZ) -> L1 | ||
| 788 | +} | ||
| 789 | +``` | ||
| 790 | + | ||
| 791 | +### A5 only — SIMT executor for Elem GM → L1 (`GatherExec::Simt`) | ||
| 792 | + | ||
| 793 | +On A5 the Elem GM → L1 path has two executors, selected by a third template parameter | ||
| 794 | +`GatherExec` (defined alongside `Coalesce`); the existing scalar path is unchanged and remains | ||
| 795 | +the default for the 4-operand form: | ||
| 796 | + | ||
| 797 | +```cpp | ||
| 798 | +enum class GatherExec : uint8_t { Scalar = 0, Simt = 1 }; | ||
| 799 | +``` | ||
| 800 | + | ||
| 801 | +- **`GatherExec::Scalar`** (the default for the 4-operand form) — the cube core walks the indices | ||
| 802 | + with scalar loads. Best for small / sparse tiles. | ||
| 803 | +- **`GatherExec::Simt`** — the **AIV vector core** collects the discrete elements with a SIMT | ||
| 804 | + kernel (`simt_mgather_l1_elem_kernel`), which parallelizes the gather across warps. This is the | ||
| 805 | + A5-only "GM → GM" stage; A2/A3 has no SIMT engine and therefore no `Simt` executor. | ||
| 806 | + | ||
| 807 | +The two executors share the **same 4-operand signature**; only the third template parameter | ||
| 808 | +selects the SIMT path. There is **no UB operand**: the SIMT engine moves data | ||
| 809 | +GM → D-cache → registers → GM, so the D-cache is implicit hardware plumbing that neither the | ||
| 810 | +caller nor the kernel manages. At the API level the SIMT executor is a pure GM → GM gather: | ||
| 811 | + | ||
| 812 | +```cpp | ||
| 813 | +template <Coalesce CMode, GatherOOB Oob, GatherExec Exec, typename MatTileDst, | ||
| 814 | + typename GlobalTable, typename GlobalIdx, typename GlobalScratch> | ||
| 815 | +PTO_INST RecordEvent MGATHER(MatTileDst& dst, GlobalTable& table, GlobalIdx& idx, | ||
| 816 | + GlobalScratch& scratch); | ||
| 817 | +``` | ||
| 818 | + | ||
| 819 | +**Algorithm.** One AIV subcore (`get_subblockid() == 0`) launches the SIMT grid over the full | ||
| 820 | +padded `Rows × Cols` NZ tile. Each thread maps its linear NZ offset back to `(r, c)` | ||
| 821 | +(`blockCol = off / (Rows * kC0)`, `r = (off % (Rows * kC0)) / kC0`, | ||
| 822 | +`c = blockCol * kC0 + (off % kC0)`), gathers `table[remap(idx[r, c])]` for in-bounds lanes and | ||
| 823 | +writes `0` for padding / `Zero`-policy lanes — so padding and OOB are handled in one pass with no | ||
| 824 | +separate pre-zero. Each thread stores its result **directly to GM `scratch`** at its linear NZ | ||
| 825 | +offset; because consecutive lanes hold consecutive offsets, the stores coalesce into contiguous | ||
| 826 | +GM bursts through the D-cache. After the grid retires, the AIV flushes the scratch range | ||
| 827 | +(`dcci` per cache line + `dsb(DSB_DDR)`) so the cube core's DMA observes the writes, then hands | ||
| 828 | +off to the AIC with a single vec → cube intra-block flag (producer `set_intra_block(PIPE_S, id)`, | ||
| 829 | +consumer `wait_intra_block(PIPE_MTE2, id)`, hardware auto-maps the subcore offset). The AIC then | ||
| 830 | +issues the contiguous `copy_gm_to_cbuf` scratch → L1. No `pipe_barrier(PIPE_ALL)` is needed in this case. | ||
| 831 | + | ||
| 832 | +**Coalescing / mapping.** The grid is sized `dim3{32, kLaunchWarps}` with | ||
| 833 | +`kLaunchWarps = min(ceil(Rows * Cols / 32), 32)`, so small tiles do not pay for idle warps. The | ||
| 834 | +linear-NZ-offset mapping makes the **scratch stores fully coalesced** (lane `i` writes offset | ||
| 835 | +`base + i`) and the **index reads contiguous within each `kC0` block**; the only unavoidable | ||
| 836 | +random traffic is `table[remap(idx)]`, which is the intrinsic cost of a gather and is serviced | ||
| 837 | +through the D-cache (warm lines are reused across lanes that hit the same row). | ||
| 838 | + | ||
| 839 | +**Constraints.** Identical to the scalar Elem path (index / table / scratch / destination); the | ||
| 840 | +SIMT executor adds no new operand-level constraints. All Elem dtypes and all `GatherOOB` policies | ||
| 841 | +are supported. | ||
| 842 | + | ||
| 843 | +```cpp | ||
| 844 | +template <typename T, int R, int C, int TableSize> | ||
| 845 | +AICORE void example_gm2l1_elem_simt(__gm__ T* tablePtr, __gm__ int32_t* idxPtr, __gm__ T* scratchPtr) | ||
| 846 | +{ | ||
| 847 | + using TableShape = Shape<1, 1, 1, 1, TableSize>; | ||
| 848 | + using TableStride = Stride<1, 1, 1, TableSize, 1>; | ||
| 849 | + using IdxShape = Shape<1, 1, 1, R, C>; | ||
| 850 | + using IdxStride = Stride<1, 1, 1, C, 1>; | ||
| 851 | + using ScratchShape = Shape<1, 1, 1, 1, R * C>; | ||
| 852 | + using ScratchStride= Stride<1, 1, 1, R * C, 1>; | ||
| 853 | + GlobalTensor<T, TableShape, TableStride, Layout::ND> tableGM(tablePtr); | ||
| 854 | + GlobalTensor<int32_t, IdxShape, IdxStride, Layout::ND> idxGM(idxPtr); | ||
| 855 | + GlobalTensor<T, ScratchShape, ScratchStride, Layout::ND> scratchGM(scratchPtr); | ||
| 856 | + | ||
| 857 | + using DstTile = Tile<TileType::Mat, T, R, C, BLayout::ColMajor, R, C, SLayout::RowMajor, 512>; | ||
| 858 | + DstTile dst; TASSIGN(dst, 0x0); | ||
| 859 | + | ||
| 860 | + // AIV SIMT gather (GM -> D-cache -> regs -> GM scratch, NZ) -> AIC copy_gm_to_cbuf -> L1. | ||
| 861 | + // No UB operand: the D-cache is implicit; the API deals only GM -> GM gather. | ||
| 862 | + MGATHER<Coalesce::Elem, GatherOOB::Zero, GatherExec::Simt>(dst, tableGM, idxGM, scratchGM); | ||
| 863 | +} | ||
| 864 | +``` | ||
| 865 | + | ||
| 651 | ## Related Instructions | 866 | ## Related Instructions |
| 652 | 867 | ||
| 653 | - [`TLOAD`](TLOAD.md): contiguous block transfer GM → Tile. | 868 | - [`TLOAD`](TLOAD.md): contiguous block transfer GM → Tile. |
| @@ -1974,6 +1974,31 @@ PTO_INST RecordEvent MGATHER(TileDst &dst, GlobalData &src, TileInd &indexes, Wa | |||
| 1974 | return {}; | 1974 | return {}; |
| 1975 | } | 1975 | } |
| 1976 | 1976 | ||
| 1977 | +template <Coalesce CMode, typename TileDst, typename GlobalData, typename GlobalIdx, typename GlobalScratch> | ||
| 1978 | +PTO_INST RecordEvent MGATHER(TileDst &dst, GlobalData &src, GlobalIdx &indexes, GlobalScratch &scratch) | ||
| 1979 | +{ | ||
| 1980 | + MGATHER_IMPL<CMode>(dst, src, indexes, scratch); | ||
| 1981 | + return {}; | ||
| 1982 | +} | ||
| 1983 | + | ||
| 1984 | +template <Coalesce CMode, GatherOOB Mode, typename TileDst, typename GlobalData, typename GlobalIdx, | ||
| 1985 | + typename GlobalScratch> | ||
| 1986 | +PTO_INST RecordEvent MGATHER(TileDst &dst, GlobalData &src, GlobalIdx &indexes, GlobalScratch &scratch) | ||
| 1987 | +{ | ||
| 1988 | + MGATHER_IMPL<CMode, Mode>(dst, src, indexes, scratch); | ||
| 1989 | + return {}; | ||
| 1990 | +} | ||
| 1991 | + | ||
| 1992 | + | ||
| 1993 | +template <Coalesce CMode, GatherOOB Mode, GatherExec Exec, typename TileDst, typename GlobalData, typename GlobalIdx, | ||
| 1994 | + typename GlobalScratch> | ||
| 1995 | +PTO_INST RecordEvent MGATHER(TileDst &dst, GlobalData &src, GlobalIdx &indexes, GlobalScratch &scratch) | ||
| 1996 | +{ | ||
| 1997 | + MGATHER_IMPL<CMode, Mode, Exec>(dst, src, indexes, scratch); | ||
| 1998 | + return {}; | ||
| 1999 | +} | ||
| 2000 | + | ||
| 2001 | + | ||
| 1977 | template <typename GlobalData, typename TileSrc, typename TileInd, typename... WaitEvents> | 2002 | template <typename GlobalData, typename TileSrc, typename TileInd, typename... WaitEvents> |
| 1978 | PTO_INST RecordEvent MSCATTER(GlobalData &dst, TileSrc &src, TileInd &indexes, WaitEvents &...events) | 2003 | PTO_INST RecordEvent MSCATTER(GlobalData &dst, TileSrc &src, TileInd &indexes, WaitEvents &...events) |
| 1979 | { | 2004 | { |
| @@ -13,6 +13,8 @@ See LICENSE in the root of the software repository for the full text of the Lice | |||
| 13 | 13 | ||
| 14 | 14 | ||
| 15 | 15 | ||
| 16 | + | ||
| 17 | + | ||
| 16 | 18 | ||
| 17 | 19 | ||
| 18 | namespace pto { | 20 | namespace pto { |
| @@ -30,30 +32,30 @@ struct IsMGatherNDTile { | |||
| 30 | static constexpr bool value = Tile::isRowMajor && (Tile::SFractal == SLayout::NoneBox); | 32 | static constexpr bool value = Tile::isRowMajor && (Tile::SFractal == SLayout::NoneBox); |
| 31 | }; | 33 | }; |
| 32 | 34 | ||
| 33 | -template <typename Tile> | ||
| 34 | -struct IsMGatherNZTile { | ||
| 35 | - static constexpr bool value = | ||
| 36 | - !Tile::isRowMajor && (Tile::SFractal == SLayout::RowMajor) && (Tile::SFractalSize == TileConfig::fractalABSize); | ||
| 37 | -}; | ||
| 38 | - | ||
| 39 | template <GatherOOB Oob> | 35 | template <GatherOOB Oob> |
| 40 | AICORE PTO_INLINE uint32_t mgather_remap(uint32_t idx, uint32_t cap, uint32_t &doRead) | 36 | AICORE PTO_INLINE uint32_t mgather_remap(uint32_t idx, uint32_t cap, uint32_t &doRead) |
| 41 | { | 37 | { |
| 42 | - if constexpr (Oob == GatherOOB::Undefined) { | 38 | + if constexpr (Oob == GatherOOB::Clamp) { |
| 43 | - doRead = 1u; | ||
| 44 | - return idx; | ||
| 45 | - } else if constexpr (Oob == GatherOOB::Clamp) { | ||
| 46 | doRead = 1u; | 39 | doRead = 1u; |
| 47 | return (idx >= cap) ? (cap - 1u) : idx; | 40 | return (idx >= cap) ? (cap - 1u) : idx; |
| 48 | } else if constexpr (Oob == GatherOOB::Wrap) { | 41 | } else if constexpr (Oob == GatherOOB::Wrap) { |
| 49 | doRead = 1u; | 42 | doRead = 1u; |
| 50 | return idx % cap; | 43 | return idx % cap; |
| 44 | + } else if constexpr (Oob == GatherOOB::Undefined) { | ||
| 45 | + doRead = 1u; | ||
| 46 | + return idx; | ||
| 51 | } else { | 47 | } else { |
| 52 | doRead = (idx < cap) ? 1u : 0u; | 48 | doRead = (idx < cap) ? 1u : 0u; |
| 53 | return idx; | 49 | return idx; |
| 54 | } | 50 | } |
| 55 | } | 51 | } |
| 56 | 52 | ||
| 53 | +template <typename Tile> | ||
| 54 | +struct IsMGatherNZTile { | ||
| 55 | + static constexpr bool value = | ||
| 56 | + !Tile::isRowMajor && (Tile::SFractal == SLayout::RowMajor) && (Tile::SFractalSize == TileConfig::fractalABSize); | ||
| 57 | +}; | ||
| 58 | + | ||
| 57 | template <typename T> | 59 | template <typename T> |
| 58 | AICORE PTO_INLINE void MGatherRowDma(__ubuf__ T *dst, __gm__ T *src, uint32_t lenBytes) | 60 | AICORE PTO_INLINE void MGatherRowDma(__ubuf__ T *dst, __gm__ T *src, uint32_t lenBytes) |
| 59 | { | 61 | { |
| @@ -84,13 +86,13 @@ AICORE PTO_INLINE uint64_t MGatherNZGmOffset(uint32_t logicalRow, uint32_t logic | |||
| 84 | int gStride0, int gStride1, int gStride2, int gStride3, int gStride4) | 86 | int gStride0, int gStride1, int gStride2, int gStride3, int gStride4) |
| 85 | { | 87 | { |
| 86 | constexpr uint32_t kC0 = C0_SIZE_BYTE / sizeof(T); | 88 | constexpr uint32_t kC0 = C0_SIZE_BYTE / sizeof(T); |
| 87 | - constexpr uint32_t kFRow = FRACTAL_NZ_ROW; | ||
| 88 | const uint32_t blockColCombined = logicalCol / kC0; | 89 | const uint32_t blockColCombined = logicalCol / kC0; |
| 89 | const uint32_t colInBlock = logicalCol - blockColCombined * kC0; | 90 | const uint32_t colInBlock = logicalCol - blockColCombined * kC0; |
| 91 | + constexpr uint32_t kFRow = FRACTAL_NZ_ROW; | ||
| 90 | const uint32_t blockRow = logicalRow / kFRow; | 92 | const uint32_t blockRow = logicalRow / kFRow; |
| 91 | - const uint32_t rowInBlock = logicalRow - blockRow * kFRow; | ||
| 92 | const uint32_t blockColOuter0 = (gShape0 == 1) ? 0u : (blockColCombined / (uint32_t)gShape1); | 93 | const uint32_t blockColOuter0 = (gShape0 == 1) ? 0u : (blockColCombined / (uint32_t)gShape1); |
| 93 | const uint32_t blockColOuter1 = (gShape0 == 1) ? blockColCombined : (blockColCombined - blockColOuter0 * gShape1); | 94 | const uint32_t blockColOuter1 = (gShape0 == 1) ? blockColCombined : (blockColCombined - blockColOuter0 * gShape1); |
| 95 | + const uint32_t rowInBlock = logicalRow - blockRow * kFRow; | ||
| 94 | return (uint64_t)blockColOuter0 * (uint64_t)gStride0 + (uint64_t)blockColOuter1 * (uint64_t)gStride1 + | 96 | return (uint64_t)blockColOuter0 * (uint64_t)gStride0 + (uint64_t)blockColOuter1 * (uint64_t)gStride1 + |
| 95 | (uint64_t)blockRow * (uint64_t)gStride2 + (uint64_t)rowInBlock * (uint64_t)gStride3 + | 97 | (uint64_t)blockRow * (uint64_t)gStride2 + (uint64_t)rowInBlock * (uint64_t)gStride3 + |
| 96 | (uint64_t)colInBlock * (uint64_t)gStride4; | 98 | (uint64_t)colInBlock * (uint64_t)gStride4; |
| @@ -288,6 +290,110 @@ __tf__ AICORE void MGatherElemNzImpl(typename DstTile::TileDType __out__ dst, __ | |||
| 288 | PtoSetWaitFlag<PIPE_S, PIPE_MTE3>(); | 290 | PtoSetWaitFlag<PIPE_S, PIPE_MTE3>(); |
| 289 | } | 291 | } |
| 290 | 292 | ||
| 293 | +template <GatherOOB Oob, typename T, typename TIdx, typename DstTile> | ||
| 294 | +__tf__ AICORE void MGatherGm2L1RowImpl(typename DstTile::TileDType __out__ dst, __gm__ T *tablePtr, __gm__ TIdx *idxPtr, | ||
| 295 | + uint32_t validRow, uint32_t validCol, uint32_t tableRows, | ||
| 296 | + uint32_t tableRowStride) | ||
| 297 | +{ | ||
| 298 | + | ||
| 299 | + constexpr uint32_t kC0 = C0_SIZE_BYTE / sizeof(T); | ||
| 300 | + __cbuf__ T *dstPtr = (__cbuf__ T *)__cce_get_tile_ptr(dst); | ||
| 301 | + constexpr uint32_t kTileRows = DstTile::Rows; | ||
| 302 | + constexpr uint32_t kTileCols = DstTile::Cols; | ||
| 303 | + | ||
| 304 | + if constexpr (Oob == GatherOOB::Zero) { | ||
| 305 | + constexpr uint32_t kColBlocks = kTileCols / kC0; | ||
| 306 | + int64_t repeatConfig = (static_cast<int64_t>(kTileRows) << 16) | static_cast<int64_t>(kColBlocks); | ||
| 307 | + pto_create_cbuf_matrix((__cbuf__ uint16_t *)dstPtr, repeatConfig, 0); | ||
| 308 | + } | ||
| 309 | + | ||
| 310 | + for (uint32_t r = 0; r < validRow; r++) { | ||
| 311 | + uint32_t doRead; | ||
| 312 | + uint32_t rawIdx = static_cast<uint32_t>(idxPtr[r]); | ||
| 313 | + uint32_t safeIdx = mgather_remap<Oob>(rawIdx, tableRows, doRead); | ||
| 314 | + if (doRead) { | ||
| 315 | + __gm__ T *srcRow = tablePtr + static_cast<uint64_t>(safeIdx) * tableRowStride; | ||
| 316 | + __cbuf__ T *dstRow = dstPtr + static_cast<uint64_t>(r) * kC0; | ||
| 317 | + pto_copy_gm_to_cbuf_multi_nd2nz<T>(dstRow, srcRow, 0, 1, 1, static_cast<uint16_t>(validCol), 0, | ||
| 318 | + static_cast<uint16_t>(tableRowStride), static_cast<uint16_t>(kTileRows), | ||
| 319 | + 1, 1); | ||
| 320 | + } | ||
| 321 | + } | ||
| 322 | + PtoSetWaitFlag<PIPE_S, PIPE_MTE2>(); | ||
| 323 | + | ||
| 324 | +} | ||
| 325 | + | ||
| 326 | +template <GatherOOB Oob, typename T, typename TIdx, typename DstTile> | ||
| 327 | +__tf__ AICORE void MGatherGm2L1ElemImpl(typename DstTile::TileDType __out__ dst, __gm__ T *tablePtr, | ||
| 328 | + __gm__ TIdx *idxPtr, __gm__ T *scratchPtr, uint32_t validRow, uint32_t validCol, | ||
| 329 | + uint32_t tableSize, uint32_t idxRowStride) | ||
| 330 | +{ | ||
| 331 | + | ||
| 332 | + constexpr uint32_t kC0 = C0_SIZE_BYTE / sizeof(T); | ||
| 333 | + constexpr uint32_t kTileCols = DstTile::Cols; | ||
| 334 | + constexpr uint32_t kTileRows = DstTile::Rows; | ||
| 335 | + constexpr uint32_t kTileNumel = kTileRows * kTileCols; | ||
| 336 | + __cbuf__ T *dstPtr = (__cbuf__ T *)__cce_get_tile_ptr(dst); | ||
| 337 | + | ||
| 338 | + for (uint32_t i = 0; i < kTileNumel; i++) { | ||
| 339 | + scratchPtr[i] = static_cast<T>(0); | ||
| 340 | + } | ||
| 341 | + for (uint32_t r = 0; r < validRow; r++) { | ||
| 342 | + const uint32_t idxRowOff = r * idxRowStride; | ||
| 343 | + for (uint32_t c = 0; c < validCol; c++) { | ||
| 344 | + uint32_t doRead; | ||
| 345 | + uint32_t rawIdx = static_cast<uint32_t>(idxPtr[idxRowOff + c]); | ||
| 346 | + uint32_t safeIdx = mgather_remap<Oob>(rawIdx, tableSize, doRead); | ||
| 347 | + if (doRead) { | ||
| 348 | + const uint32_t blockCol = c / kC0; | ||
| 349 | + const uint32_t colInBlock = c - blockCol * kC0; | ||
| 350 | + const uint64_t off = | ||
| 351 | + static_cast<uint64_t>(blockCol) * static_cast<uint64_t>(kTileRows) * static_cast<uint64_t>(kC0) + | ||
| 352 | + static_cast<uint64_t>(r) * static_cast<uint64_t>(kC0) + static_cast<uint64_t>(colInBlock); | ||
| 353 | + scratchPtr[off] = tablePtr[safeIdx]; | ||
| 354 | + } | ||
| 355 | + } | ||
| 356 | + } | ||
| 357 | + const uint32_t totalBytes = kTileNumel * sizeof(T); | ||
| 358 | + constexpr uint32_t kCacheLineBytes = 64; | ||
| 359 | + const uint32_t numLines = (totalBytes + kCacheLineBytes - 1) / kCacheLineBytes; | ||
| 360 | + __gm__ uint8_t *flushPtr = reinterpret_cast<__gm__ uint8_t *>(scratchPtr); | ||
| 361 | + for (uint32_t i = 0; i < numLines; i++) { | ||
| 362 | + dcci(static_cast<__gm__ void *>(flushPtr + i * kCacheLineBytes), SINGLE_CACHE_LINE); | ||
| 363 | + } | ||
| 364 | + dsb(DSB_DDR); | ||
| 365 | + PtoSetWaitFlag<PIPE_S, PIPE_MTE2>(); | ||
| 366 | + const uint16_t lenBurst = static_cast<uint16_t>(kTileNumel * sizeof(T) / BLOCK_BYTE_SIZE); | ||
| 367 | + copy_gm_to_cbuf(dstPtr, scratchPtr, (uint8_t)0, (uint16_t)1, lenBurst, (uint16_t)0, (uint16_t)0, (pad_t)0); | ||
| 368 | + | ||
| 369 | +} | ||
| 370 | + | ||
| 371 | +template <Coalesce Mode, GatherOOB Oob, typename DstTile, typename GlobalTable, typename IdxSrc> | ||
| 372 | +PTO_INTERNAL void MGatherCheckGm2L1() | ||
| 373 | +{ | ||
| 374 | + using T = typename DstTile::DType; | ||
| 375 | + | ||
| 376 | + static_assert(IsValidMGatherDType<T>::value, | ||
| 377 | + "MGATHER A2/A3 GM->L1 data type must be int8/uint8/int16/uint16/int32/uint32/half/bfloat16/float."); | ||
| 378 | + static_assert(std::is_same_v<typename GlobalTable::DType, __gm__ T>, | ||
| 379 | + "MGATHER A2/A3 GM->L1 table must be a GM GlobalTensor with element type matching the destination."); | ||
| 380 | + static_assert(std::is_same_v<typename IdxSrc::DType, __gm__ int32_t> || | ||
| 381 | + std::is_same_v<typename IdxSrc::DType, __gm__ uint32_t>, | ||
| 382 | + "MGATHER A2/A3 GM->L1 indices must be a GM int32_t/uint32_t GlobalTensor."); | ||
| 383 | + static_assert(DstTile::Loc == TileType::Mat, "MGATHER A2/A3 GM->L1 destination must be a Mat tile (L1)."); | ||
| 384 | + static_assert(IsMGatherNZTile<DstTile>::value, | ||
| 385 | + "MGATHER A2/A3 GM->L1 destination must be NZ (BLayout::ColMajor + SLayout::RowMajor + " | ||
| 386 | + "SFractalSize=512)."); | ||
| 387 | + static_assert(GlobalTable::layout == Layout::ND, "MGATHER A2/A3 GM->L1 table must use Layout::ND."); | ||
| 388 | + static_assert(DstTile::Cols % (C0_SIZE_BYTE / sizeof(T)) == 0, | ||
| 389 | + "MGATHER A2/A3 GM->L1 destination tile Cols must be a multiple of C0 (= 32 / sizeof(T))."); | ||
| 390 | + static_assert(DstTile::Rows % FRACTAL_NZ_ROW == 0, | ||
| 391 | + "MGATHER A2/A3 GM->L1 destination tile Rows must be a multiple of FRACTAL_NZ_ROW (16)."); | ||
| 392 | + if constexpr (Mode == Coalesce::Row) { | ||
| 393 | + static_assert(sizeof(T) <= 4, "MGATHER A2/A3 GM->L1 Coalesce::Row supports b8/b16/b32 element types."); | ||
| 394 | + } | ||
| 395 | +} | ||
| 396 | + | ||
| 291 | template <Coalesce Mode, GatherOOB Oob, typename DstTile, typename GlobalTable, typename IdxTile> | 397 | template <Coalesce Mode, GatherOOB Oob, typename DstTile, typename GlobalTable, typename IdxTile> |
| 292 | PTO_INTERNAL void MGatherCheck() | 398 | PTO_INTERNAL void MGatherCheck() |
| 293 | { | 399 | { |
| @@ -358,55 +464,100 @@ template <Coalesce Mode = Coalesce::Row, GatherOOB Oob = GatherOOB::Undefined, t | |||
| 358 | PTO_INTERNAL void MGATHER_IMPL(DstTile &dst, GlobalTable &table, IdxTile &indices) | 464 | PTO_INTERNAL void MGATHER_IMPL(DstTile &dst, GlobalTable &table, IdxTile &indices) |
| 359 | { | 465 | { |
| 360 | using T = typename DstTile::DType; | 466 | using T = typename DstTile::DType; |
| 361 | - using TIdx = typename IdxTile::DType; | ||
| 362 | 467 | ||
| 363 | - MGatherCheck<Mode, Oob, DstTile, GlobalTable, IdxTile>(); | 468 | + if constexpr (DstTile::Loc == TileType::Mat) { |
| 469 | + MGatherCheckGm2L1<Coalesce::Row, Oob, DstTile, GlobalTable, IdxTile>(); | ||
| 470 | + using TIdx = std::conditional_t<std::is_same_v<typename IdxTile::DType, __gm__ uint32_t>, uint32_t, int32_t>; | ||
| 471 | + __gm__ T *tablePtr = reinterpret_cast<__gm__ T *>(table.data()); | ||
| 472 | + __gm__ TIdx *idxPtr = reinterpret_cast<__gm__ TIdx *>(indices.data()); | ||
| 473 | + const uint32_t validRow = dst.GetValidRow(); | ||
| 474 | + const uint32_t validCol = dst.GetValidCol(); | ||
| 475 | + const uint32_t tableRows = | ||
| 476 | + static_cast<uint32_t>(table.GetShape(GlobalTensorDim::DIM_0) * table.GetShape(GlobalTensorDim::DIM_1) * | ||
| 477 | + table.GetShape(GlobalTensorDim::DIM_2) * table.GetShape(GlobalTensorDim::DIM_3)); | ||
| 478 | + const uint32_t tableRowStride = static_cast<uint32_t>(table.GetStride(GlobalTensorDim::DIM_3)); | ||
| 479 | + MGatherGm2L1RowImpl<Oob, T, TIdx, DstTile>(dst.data(), tablePtr, idxPtr, validRow, validCol, tableRows, | ||
| 480 | + tableRowStride); | ||
| 481 | + return; | ||
| 482 | + } else { | ||
| 483 | + using TIdx = typename IdxTile::DType; | ||
| 364 | 484 | ||
| 485 | + MGatherCheck<Mode, Oob, DstTile, GlobalTable, IdxTile>(); | ||
| 486 | + | ||
| 487 | + __gm__ T *tablePtr = reinterpret_cast<__gm__ T *>(table.data()); | ||
| 488 | + | ||
| 489 | + const uint32_t validRow = dst.GetValidRow(); | ||
| 490 | + const uint32_t validCol = dst.GetValidCol(); | ||
| 491 | + | ||
| 492 | + constexpr bool kIsTableNZ = (GlobalTable::layout == Layout::NZ); | ||
| 493 | + | ||
| 494 | + if constexpr (kIsTableNZ) { | ||
| 495 | + const int gShape0 = static_cast<int>(table.GetShape(GlobalTensorDim::DIM_0)); | ||
| 496 | + const int gShape1 = static_cast<int>(table.GetShape(GlobalTensorDim::DIM_1)); | ||
| 497 | + const int gShape2 = static_cast<int>(table.GetShape(GlobalTensorDim::DIM_2)); | ||
| 498 | + const int gStride0 = static_cast<int>(table.GetStride(GlobalTensorDim::DIM_0)); | ||
| 499 | + const int gStride1 = static_cast<int>(table.GetStride(GlobalTensorDim::DIM_1)); | ||
| 500 | + const int gStride2 = static_cast<int>(table.GetStride(GlobalTensorDim::DIM_2)); | ||
| 501 | + const int gStride3 = static_cast<int>(table.GetStride(GlobalTensorDim::DIM_3)); | ||
| 502 | + const int gStride4 = static_cast<int>(table.GetStride(GlobalTensorDim::DIM_4)); | ||
| 503 | + | ||
| 504 | + if constexpr (Mode == Coalesce::Row) { | ||
| 505 | + MGatherRowNzImpl<Oob, T, TIdx, DstTile, IdxTile>(dst.data(), tablePtr, indices.data(), validRow, | ||
| 506 | + gShape0, gShape1, gShape2, gStride0, gStride1, | ||
| 507 | + gStride2, gStride3); | ||
| 508 | + } else { | ||
| 509 | + constexpr uint32_t kC0 = C0_SIZE_BYTE / sizeof(T); | ||
| 510 | + const uint32_t nLogicalCols = static_cast<uint32_t>(gShape0 * gShape1) * kC0; | ||
| 511 | + const uint32_t tableSize = static_cast<uint32_t>(gShape2 * FRACTAL_NZ_ROW) * nLogicalCols; | ||
| 512 | + MGatherElemNzImpl<Oob, T, TIdx, DstTile, IdxTile>(dst.data(), tablePtr, indices.data(), validRow, | ||
| 513 | + validCol, tableSize, gShape0, gShape1, gStride0, | ||
| 514 | + gStride1, gStride2, gStride3, gStride4, nLogicalCols); | ||
| 515 | + } | ||
| 516 | + } else { | ||
| 517 | + if constexpr (Mode == Coalesce::Row) { | ||
| 518 | + const uint32_t tableRows = static_cast<uint32_t>( | ||
| 519 | + table.GetShape(GlobalTensorDim::DIM_0) * table.GetShape(GlobalTensorDim::DIM_1) * | ||
| 520 | + table.GetShape(GlobalTensorDim::DIM_2) * table.GetShape(GlobalTensorDim::DIM_3)); | ||
| 521 | + const uint32_t tableRowStride = static_cast<uint32_t>(table.GetStride(GlobalTensorDim::DIM_3)); | ||
| 522 | + MGatherRowImpl<Oob, T, TIdx, DstTile, IdxTile>(dst.data(), tablePtr, indices.data(), validRow, validCol, | ||
| 523 | + tableRows, tableRowStride); | ||
| 524 | + } else { | ||
| 525 | + const uint32_t tableSize = static_cast<uint32_t>( | ||
| 526 | + table.GetShape(GlobalTensorDim::DIM_0) * table.GetShape(GlobalTensorDim::DIM_1) * | ||
| 527 | + table.GetShape(GlobalTensorDim::DIM_2) * table.GetShape(GlobalTensorDim::DIM_3) * | ||
| 528 | + table.GetShape(GlobalTensorDim::DIM_4)); | ||
| 529 | + MGatherElemImpl<Oob, T, TIdx, DstTile, IdxTile>(dst.data(), tablePtr, indices.data(), validRow, | ||
| 530 | + validCol, tableSize); | ||
| 531 | + } | ||
| 532 | + } | ||
| 533 | + } | ||
| 534 | +} | ||
| 535 | + | ||
| 536 | +template <Coalesce Mode = Coalesce::Elem, GatherOOB Oob = GatherOOB::Undefined, typename DstTile, typename GlobalTable, | ||
| 537 | + typename IdxSrc, typename GlobalScratch> | ||
| 538 | +PTO_INTERNAL void MGATHER_IMPL(DstTile &dst, GlobalTable &table, IdxSrc &indices, GlobalScratch &scratch) | ||
| 539 | +{ | ||
| 540 | + using T = typename DstTile::DType; | ||
| 541 | + | ||
| 542 | + MGatherCheckGm2L1<Coalesce::Elem, Oob, DstTile, GlobalTable, IdxSrc>(); | ||
| 543 | + static_assert(std::is_same_v<typename GlobalScratch::DType, __gm__ T>, | ||
| 544 | + "MGATHER A2/A3 GM->L1 scratch must be a GM GlobalTensor with element type matching the destination."); | ||
| 545 | + | ||
| 546 | + using TIdx = std::conditional_t<std::is_same_v<typename IdxSrc::DType, __gm__ uint32_t>, uint32_t, int32_t>; | ||
| 365 | __gm__ T *tablePtr = reinterpret_cast<__gm__ T *>(table.data()); | 547 | __gm__ T *tablePtr = reinterpret_cast<__gm__ T *>(table.data()); |
| 548 | + __gm__ TIdx *idxPtr = reinterpret_cast<__gm__ TIdx *>(indices.data()); | ||
| 549 | + __gm__ T *scratchPtr = reinterpret_cast<__gm__ T *>(scratch.data()); | ||
| 366 | 550 | ||
| 367 | const uint32_t validRow = dst.GetValidRow(); | 551 | const uint32_t validRow = dst.GetValidRow(); |
| 368 | const uint32_t validCol = dst.GetValidCol(); | 552 | const uint32_t validCol = dst.GetValidCol(); |
| 553 | + const uint32_t tableSize = | ||
| 554 | + static_cast<uint32_t>(table.GetShape(GlobalTensorDim::DIM_0) * table.GetShape(GlobalTensorDim::DIM_1) * | ||
| 555 | + table.GetShape(GlobalTensorDim::DIM_2) * table.GetShape(GlobalTensorDim::DIM_3) * | ||
| 556 | + table.GetShape(GlobalTensorDim::DIM_4)); | ||
| 557 | + const uint32_t idxRowStride = static_cast<uint32_t>(indices.GetStride(GlobalTensorDim::DIM_3)); | ||
| 369 | 558 | ||
| 370 | - constexpr bool kIsTableNZ = (GlobalTable::layout == Layout::NZ); | 559 | + MGatherGm2L1ElemImpl<Oob, T, TIdx, DstTile>(dst.data(), tablePtr, idxPtr, scratchPtr, validRow, validCol, tableSize, |
| 371 | - | 560 | + idxRowStride); |
| 372 | - if constexpr (kIsTableNZ) { | ||
| 373 | - const int gShape0 = static_cast<int>(table.GetShape(GlobalTensorDim::DIM_0)); | ||
| 374 | - const int gShape1 = static_cast<int>(table.GetShape(GlobalTensorDim::DIM_1)); | ||
| 375 | - const int gShape2 = static_cast<int>(table.GetShape(GlobalTensorDim::DIM_2)); | ||
| 376 | - const int gStride0 = static_cast<int>(table.GetStride(GlobalTensorDim::DIM_0)); | ||
| 377 | - const int gStride1 = static_cast<int>(table.GetStride(GlobalTensorDim::DIM_1)); | ||
| 378 | - const int gStride2 = static_cast<int>(table.GetStride(GlobalTensorDim::DIM_2)); | ||
| 379 | - const int gStride3 = static_cast<int>(table.GetStride(GlobalTensorDim::DIM_3)); | ||
| 380 | - const int gStride4 = static_cast<int>(table.GetStride(GlobalTensorDim::DIM_4)); | ||
| 381 | - | ||
| 382 | - if constexpr (Mode == Coalesce::Row) { | ||
| 383 | - MGatherRowNzImpl<Oob, T, TIdx, DstTile, IdxTile>(dst.data(), tablePtr, indices.data(), validRow, gShape0, | ||
| 384 | - gShape1, gShape2, gStride0, gStride1, gStride2, gStride3); | ||
| 385 | - } else { | ||
| 386 | - constexpr uint32_t kC0 = C0_SIZE_BYTE / sizeof(T); | ||
| 387 | - const uint32_t nLogicalCols = static_cast<uint32_t>(gShape0 * gShape1) * kC0; | ||
| 388 | - const uint32_t tableSize = static_cast<uint32_t>(gShape2 * FRACTAL_NZ_ROW) * nLogicalCols; | ||
| 389 | - MGatherElemNzImpl<Oob, T, TIdx, DstTile, IdxTile>(dst.data(), tablePtr, indices.data(), validRow, validCol, | ||
| 390 | - tableSize, gShape0, gShape1, gStride0, gStride1, gStride2, | ||
| 391 | - gStride3, gStride4, nLogicalCols); | ||
| 392 | - } | ||
| 393 | - } else { | ||
| 394 | - if constexpr (Mode == Coalesce::Row) { | ||
| 395 | - const uint32_t tableRows = | ||
| 396 | - static_cast<uint32_t>(table.GetShape(GlobalTensorDim::DIM_0) * table.GetShape(GlobalTensorDim::DIM_1) * | ||
| 397 | - table.GetShape(GlobalTensorDim::DIM_2) * table.GetShape(GlobalTensorDim::DIM_3)); | ||
| 398 | - const uint32_t tableRowStride = static_cast<uint32_t>(table.GetStride(GlobalTensorDim::DIM_3)); | ||
| 399 | - MGatherRowImpl<Oob, T, TIdx, DstTile, IdxTile>(dst.data(), tablePtr, indices.data(), validRow, validCol, | ||
| 400 | - tableRows, tableRowStride); | ||
| 401 | - } else { | ||
| 402 | - const uint32_t tableSize = | ||
| 403 | - static_cast<uint32_t>(table.GetShape(GlobalTensorDim::DIM_0) * table.GetShape(GlobalTensorDim::DIM_1) * | ||
| 404 | - table.GetShape(GlobalTensorDim::DIM_2) * table.GetShape(GlobalTensorDim::DIM_3) * | ||
| 405 | - table.GetShape(GlobalTensorDim::DIM_4)); | ||
| 406 | - MGatherElemImpl<Oob, T, TIdx, DstTile, IdxTile>(dst.data(), tablePtr, indices.data(), validRow, validCol, | ||
| 407 | - tableSize); | ||
| 408 | - } | ||
| 409 | - } | ||
| 410 | } | 561 | } |
| 411 | 562 | ||
| 412 | } // namespace pto | 563 | } // namespace pto |
| @@ -13,11 +13,22 @@ See LICENSE in the root of the software repository for the full text of the Lice | |||
| 13 | 13 | ||
| 14 | 14 | ||
| 15 | 15 | ||
| 16 | + | ||
| 17 | + | ||
| 16 | 18 | ||
| 17 | 19 | ||
| 18 | 20 | ||
| 19 | namespace pto { | 21 | namespace pto { |
| 20 | 22 | ||
| 23 | + | ||
| 24 | + | ||
| 25 | +enum class GatherExec : uint8_t | ||
| 26 | +{ | ||
| 27 | + Scalar = 0, | ||
| 28 | + Simt = 1 | ||
| 29 | +}; | ||
| 30 | + | ||
| 31 | + | ||
| 21 | template <typename T> | 32 | template <typename T> |
| 22 | struct IsValidGatherDType { | 33 | struct IsValidGatherDType { |
| 23 | static constexpr bool value = | 34 | static constexpr bool value = |
| @@ -36,12 +47,12 @@ constexpr uint32_t MAX_THREADS = WARP_SIZE * MAX_WARPS; | |||
| 36 | template <GatherOOB Oob> | 47 | template <GatherOOB Oob> |
| 37 | __simt_callee__ AICORE PTO_INLINE uint32_t gather_remap(uint32_t idx, uint32_t cap, uint32_t &doRead) | 48 | __simt_callee__ AICORE PTO_INLINE uint32_t gather_remap(uint32_t idx, uint32_t cap, uint32_t &doRead) |
| 38 | { | 49 | { |
| 39 | - if constexpr (Oob == GatherOOB::Undefined) { | 50 | + if constexpr (Oob == GatherOOB::Clamp) { |
| 40 | - doRead = 1u; | ||
| 41 | - return idx; | ||
| 42 | - } else if constexpr (Oob == GatherOOB::Clamp) { | ||
| 43 | doRead = 1u; | 51 | doRead = 1u; |
| 44 | return (idx >= cap) ? (cap - 1u) : idx; | 52 | return (idx >= cap) ? (cap - 1u) : idx; |
| 53 | + } else if constexpr (Oob == GatherOOB::Undefined) { | ||
| 54 | + doRead = 1u; | ||
| 55 | + return idx; | ||
| 45 | } else if constexpr (Oob == GatherOOB::Wrap) { | 56 | } else if constexpr (Oob == GatherOOB::Wrap) { |
| 46 | doRead = 1u; | 57 | doRead = 1u; |
| 47 | return idx % cap; | 58 | return idx % cap; |
| @@ -210,6 +221,234 @@ __tf__ AICORE void MGatherScalarImpl(typename DstTileData::TileDType __out__ dst | |||
| 210 | wait_flag(PIPE_S, PIPE_V, EVENT_ID0); | 221 | wait_flag(PIPE_S, PIPE_V, EVENT_ID0); |
| 211 | } | 222 | } |
| 212 | 223 | ||
| 224 | +template <typename Tile> | ||
| 225 | +struct IsMGatherNZTile { | ||
| 226 | + static constexpr bool value = | ||
| 227 | + !Tile::isRowMajor && (Tile::SFractal == SLayout::RowMajor) && (Tile::SFractalSize == TileConfig::fractalABSize); | ||
| 228 | +}; | ||
| 229 | + | ||
| 230 | +template <GatherOOB Oob> | ||
| 231 | +AICORE PTO_INLINE uint32_t gather_remap_l1(uint32_t idx, uint32_t cap, uint32_t &doRead) | ||
| 232 | +{ | ||
| 233 | + if constexpr (Oob == GatherOOB::Undefined) { | ||
| 234 | + doRead = 1u; | ||
| 235 | + return idx; | ||
| 236 | + } else if constexpr (Oob == GatherOOB::Clamp) { | ||
| 237 | + doRead = 1u; | ||
| 238 | + return (idx >= cap) ? (cap - 1u) : idx; | ||
| 239 | + } else if constexpr (Oob == GatherOOB::Wrap) { | ||
| 240 | + doRead = 1u; | ||
| 241 | + return idx % cap; | ||
| 242 | + } else { | ||
| 243 | + doRead = (idx < cap) ? 1u : 0u; | ||
| 244 | + return idx; | ||
| 245 | + } | ||
| 246 | +} | ||
| 247 | + | ||
| 248 | +template <GatherOOB Oob, typename T, typename TIdx, typename DstTile> | ||
| 249 | +__tf__ AICORE void MGatherGm2L1RowImpl(typename DstTile::TileDType __out__ dst, __gm__ const T *tablePtr, | ||
| 250 | + __gm__ const TIdx *idxPtr, uint32_t validRow, uint32_t validCol, | ||
| 251 | + uint32_t tableRows, uint32_t tableRowStride) | ||
| 252 | +{ | ||
| 253 | + | ||
| 254 | + constexpr uint32_t kC0 = C0_SIZE_BYTE / sizeof(T); | ||
| 255 | + constexpr uint32_t kTileRows = DstTile::Rows; | ||
| 256 | + constexpr uint32_t kTileCols = DstTile::Cols; | ||
| 257 | + __cbuf__ T *dstPtr = (__cbuf__ T *)__cce_get_tile_ptr(dst); | ||
| 258 | + | ||
| 259 | + if constexpr (Oob == GatherOOB::Zero) { | ||
| 260 | + constexpr uint32_t kColBlocks = kTileCols / kC0; | ||
| 261 | + int64_t repeatConfig = (static_cast<int64_t>(kTileRows) << 16) | static_cast<int64_t>(kColBlocks); | ||
| 262 | + pto_create_cbuf_matrix((__cbuf__ uint16_t *)dstPtr, repeatConfig, 0); | ||
| 263 | + } | ||
| 264 | + | ||
| 265 | + constexpr uint16_t ndNum = 1; | ||
| 266 | + constexpr uint16_t loop2DstStride = 1; | ||
| 267 | + constexpr uint16_t loop3DstStride = kTileRows; | ||
| 268 | + constexpr uint16_t loop4DstStride = 0; | ||
| 269 | + uint64_t mte2NzPara = static_cast<uint64_t>(loop4DstStride) << 48; | ||
| 270 | + mte2NzPara |= static_cast<uint64_t>(loop3DstStride) << 32; | ||
| 271 | + mte2NzPara |= static_cast<uint64_t>(loop2DstStride) << 16; | ||
| 272 | + mte2NzPara |= static_cast<uint64_t>(ndNum); | ||
| 273 | + set_mte2_nz_para(mte2NzPara); | ||
| 274 | + | ||
| 275 | + const uint64_t loop1SrcStride = static_cast<uint64_t>(tableRowStride) * sizeof(T); | ||
| 276 | + for (uint32_t r = 0; r < validRow; r++) { | ||
| 277 | + uint32_t rawIdx = static_cast<uint32_t>(idxPtr[r]); | ||
| 278 | + uint32_t doRead; | ||
| 279 | + uint32_t safeIdx = gather_remap_l1<Oob>(rawIdx, tableRows, doRead); | ||
| 280 | + if (doRead) { | ||
| 281 | + __gm__ const T *srcRow = tablePtr + static_cast<uint64_t>(safeIdx) * tableRowStride; | ||
| 282 | + __cbuf__ T *dstRow = dstPtr + static_cast<uint64_t>(r) * kC0; | ||
| 283 | + pto_copy_gm_to_cbuf_multi_nd2nz<T>(dstRow, const_cast<__gm__ T *>(srcRow), 0, loop1SrcStride, 0, 1, | ||
| 284 | + static_cast<uint32_t>(validCol), 0); | ||
| 285 | + } | ||
| 286 | + } | ||
| 287 | + set_flag(PIPE_S, PIPE_MTE2, EVENT_ID0); | ||
| 288 | + wait_flag(PIPE_S, PIPE_MTE2, EVENT_ID0); | ||
| 289 | + | ||
| 290 | +} | ||
| 291 | + | ||
| 292 | +template <GatherOOB Oob, typename T, typename TIdx, typename DstTile> | ||
| 293 | +__tf__ AICORE void MGatherGm2L1ElemImpl(typename DstTile::TileDType __out__ dst, __gm__ const T *tablePtr, | ||
| 294 | + __gm__ const TIdx *idxPtr, __gm__ T *scratchPtr, uint32_t validRow, | ||
| 295 | + uint32_t validCol, uint32_t tableSize, uint32_t idxRowStride) | ||
| 296 | +{ | ||
| 297 | + | ||
| 298 | + __cbuf__ T *dstPtr = (__cbuf__ T *)__cce_get_tile_ptr(dst); | ||
| 299 | + constexpr uint32_t kC0 = C0_SIZE_BYTE / sizeof(T); | ||
| 300 | + constexpr uint32_t kTileRows = DstTile::Rows; | ||
| 301 | + constexpr uint32_t kTileCols = DstTile::Cols; | ||
| 302 | + constexpr uint32_t kTileNumel = kTileRows * kTileCols; | ||
| 303 | + | ||
| 304 | + for (uint32_t i = 0; i < kTileNumel; i++) { | ||
| 305 | + scratchPtr[i] = static_cast<T>(0); | ||
| 306 | + } | ||
| 307 | + for (uint32_t r = 0; r < validRow; r++) { | ||
| 308 | + const uint32_t idxRowOff = r * idxRowStride; | ||
| 309 | + for (uint32_t c = 0; c < validCol; c++) { | ||
| 310 | + uint32_t rawIdx = static_cast<uint32_t>(idxPtr[idxRowOff + c]); | ||
| 311 | + uint32_t doRead; | ||
| 312 | + uint32_t safeIdx = gather_remap_l1<Oob>(rawIdx, tableSize, doRead); | ||
| 313 | + if (doRead == 1) { | ||
| 314 | + const uint32_t blockCol = c / kC0; | ||
| 315 | + const uint32_t colInBlock = c - blockCol * kC0; | ||
| 316 | + const uint64_t off = | ||
| 317 | + static_cast<uint64_t>(blockCol) * static_cast<uint64_t>(kTileRows) * static_cast<uint64_t>(kC0) + | ||
| 318 | + static_cast<uint64_t>(r) * static_cast<uint64_t>(kC0) + static_cast<uint64_t>(colInBlock); | ||
| 319 | + scratchPtr[off] = tablePtr[safeIdx]; | ||
| 320 | + } | ||
| 321 | + } | ||
| 322 | + } | ||
| 323 | + constexpr uint32_t kCacheLineBytes = 64; | ||
| 324 | + const uint32_t totalBytes = kTileNumel * sizeof(T); | ||
| 325 | + const uint32_t numLines = (totalBytes + kCacheLineBytes - 1) / kCacheLineBytes; | ||
| 326 | + __gm__ uint8_t *flushPtr = reinterpret_cast<__gm__ uint8_t *>(scratchPtr); | ||
| 327 | + for (uint32_t i = 0; i < numLines; i++) { | ||
| 328 | + dcci(static_cast<__gm__ void *>(flushPtr + i * kCacheLineBytes), SINGLE_CACHE_LINE); | ||
| 329 | + } | ||
| 330 | + dsb(DSB_DDR); | ||
| 331 | + set_flag(PIPE_S, PIPE_MTE2, EVENT_ID0); | ||
| 332 | + wait_flag(PIPE_S, PIPE_MTE2, EVENT_ID0); | ||
| 333 | + const uint32_t lenBurst = kTileNumel * sizeof(T); | ||
| 334 | + pto_copy_gm_to_cbuf_align_v2<T>(dstPtr, scratchPtr, 0, 1, lenBurst, 0, 0, true, 0, 0, 0); | ||
| 335 | + | ||
| 336 | +} | ||
| 337 | + | ||
| 338 | +template <typename T, typename TIdx, GatherOOB Oob, uint32_t TileRowsT, uint32_t TileColsT, uint32_t ValidRowsT, | ||
| 339 | + uint32_t ValidColsT, uint32_t TableSizeT> | ||
| 340 | +AICORE __simt_vf__ LAUNCH_BOUND(1024) PTO_INLINE | ||
| 341 | + void simt_mgather_l1_elem_kernel(__gm__ T *__restrict__ scratch, __gm__ const T *__restrict__ table, | ||
| 342 | + __gm__ const TIdx *__restrict__ indices, uint32_t validRowsRT, | ||
| 343 | + uint32_t validColsRT, uint32_t tableSizeRT, uint32_t idxRowStrideRT) | ||
| 344 | +{ | ||
| 345 | + constexpr uint32_t kC0 = C0_SIZE_BYTE / sizeof(T); | ||
| 346 | + constexpr uint32_t kTileRows = TileRowsT; | ||
| 347 | + constexpr uint32_t kTileNumel = TileRowsT * TileColsT; | ||
| 348 | + constexpr uint32_t kBlockSpan = kTileRows * kC0; | ||
| 349 | + | ||
| 350 | + const uint32_t validRows = (ValidRowsT > 0u) ? ValidRowsT : validRowsRT; | ||
| 351 | + const uint32_t validCols = (ValidColsT > 0u) ? ValidColsT : validColsRT; | ||
| 352 | + const uint32_t tableSize = (TableSizeT > 0u) ? TableSizeT : tableSizeRT; | ||
| 353 | + const uint32_t idxRowStride = idxRowStrideRT; | ||
| 354 | + | ||
| 355 | + const uint32_t kNeededWarps = (kTileNumel + mgather_cfg::WARP_SIZE - 1u) / mgather_cfg::WARP_SIZE; | ||
| 356 | + const uint32_t kLaunchWarps = | ||
| 357 | + (kNeededWarps == 0u) ? 1u : ((kNeededWarps < mgather_cfg::MAX_WARPS) ? kNeededWarps : mgather_cfg::MAX_WARPS); | ||
| 358 | + const uint32_t kLaunchThreads = kLaunchWarps * mgather_cfg::WARP_SIZE; | ||
| 359 | + | ||
| 360 | + const uint32_t tid = threadIdx.y * mgather_cfg::WARP_SIZE + threadIdx.x; | ||
| 361 | + | ||
| 362 | + | ||
| 363 | + for (uint32_t off = tid; off < kTileNumel; off += kLaunchThreads) { | ||
| 364 | + const uint32_t blockCol = off / kBlockSpan; | ||
| 365 | + const uint32_t rem = off - blockCol * kBlockSpan; | ||
| 366 | + const uint32_t r = rem / kC0; | ||
| 367 | + const uint32_t colInBlock = rem - r * kC0; | ||
| 368 | + const uint32_t c = blockCol * kC0 + colInBlock; | ||
| 369 | + T val = static_cast<T>(0); | ||
| 370 | + if (r < validRows && c < validCols) { | ||
| 371 | + const uint32_t rawIdx = static_cast<uint32_t>(indices[r * idxRowStride + c]); | ||
| 372 | + uint32_t doRead; | ||
| 373 | + const uint32_t safeIdx = gather_remap<Oob>(rawIdx, tableSize, doRead); | ||
| 374 | + val = doRead ? table[safeIdx] : static_cast<T>(0); | ||
| 375 | + } | ||
| 376 | + scratch[off] = val; | ||
| 377 | + } | ||
| 378 | +} | ||
| 379 | + | ||
| 380 | +template <GatherOOB Oob, typename T, typename TIdx, typename DstTile, uint8_t SyncId> | ||
| 381 | +__tf__ AICORE void MGatherGm2L1ElemSimtImpl(typename DstTile::TileDType __out__ dst, __gm__ const T *tablePtr, | ||
| 382 | + __gm__ const TIdx *idxPtr, __gm__ T *scratchPtr, uint32_t validRow, | ||
| 383 | + uint32_t validCol, uint32_t tableSize, uint32_t idxRowStride) | ||
| 384 | +{ | ||
| 385 | + constexpr uint32_t kTileRows = DstTile::Rows; | ||
| 386 | + constexpr uint32_t kTileCols = DstTile::Cols; | ||
| 387 | + constexpr uint32_t kTileNumel = kTileRows * kTileCols; | ||
| 388 | + | ||
| 389 | + if (get_subblockid() == 0) { | ||
| 390 | + const uint32_t needed = (kTileNumel + mgather_cfg::WARP_SIZE - 1u) / mgather_cfg::WARP_SIZE; | ||
| 391 | + const uint32_t launchWarps = | ||
| 392 | + (needed == 0u) ? 1u : ((needed < mgather_cfg::MAX_WARPS) ? needed : mgather_cfg::MAX_WARPS); | ||
| 393 | + cce::async_invoke<simt_mgather_l1_elem_kernel<T, TIdx, Oob, kTileRows, kTileCols, 0u, 0u, 0u>>( | ||
| 394 | + cce::dim3{mgather_cfg::WARP_SIZE, launchWarps}, scratchPtr, tablePtr, idxPtr, validRow, validCol, tableSize, | ||
| 395 | + idxRowStride); | ||
| 396 | + | ||
| 397 | + set_flag(PIPE_V, PIPE_S, EVENT_ID0); | ||
| 398 | + wait_flag(PIPE_V, PIPE_S, EVENT_ID0); | ||
| 399 | + | ||
| 400 | + constexpr uint32_t kCacheLineBytes = 64; | ||
| 401 | + const uint32_t totalBytes = kTileNumel * sizeof(T); | ||
| 402 | + const uint32_t numLines = (totalBytes + kCacheLineBytes - 1) / kCacheLineBytes; | ||
| 403 | + __gm__ uint8_t *flushPtr = reinterpret_cast<__gm__ uint8_t *>(scratchPtr); | ||
| 404 | + for (uint32_t i = 0; i < numLines; i++) { | ||
| 405 | + dcci(static_cast<__gm__ void *>(flushPtr + i * kCacheLineBytes), SINGLE_CACHE_LINE); | ||
| 406 | + } | ||
| 407 | + dsb(DSB_DDR); | ||
| 408 | + set_intra_block(PIPE_S, SyncId); | ||
| 409 | + } | ||
| 410 | + | ||
| 411 | + | ||
| 412 | + __cbuf__ T *dstPtr = (__cbuf__ T *)__cce_get_tile_ptr(dst); | ||
| 413 | + wait_intra_block(PIPE_MTE2, SyncId); | ||
| 414 | + const uint32_t lenBurst = kTileNumel * sizeof(T); | ||
| 415 | + pto_copy_gm_to_cbuf_align_v2<T>(dstPtr, scratchPtr, 0, 1, lenBurst, 0, 0, true, 0, 0, 0); | ||
| 416 | + (void)tablePtr; | ||
| 417 | + (void)idxPtr; | ||
| 418 | + (void)validRow; | ||
| 419 | + (void)validCol; | ||
| 420 | + (void)tableSize; | ||
| 421 | + (void)idxRowStride; | ||
| 422 | + | ||
| 423 | +} | ||
| 424 | + | ||
| 425 | +template <Coalesce Mode, GatherOOB Oob, typename DstTile, typename GlobalTable, typename IdxSrc> | ||
| 426 | +PTO_INTERNAL void MGatherCheckGm2L1() | ||
| 427 | +{ | ||
| 428 | + using T = typename DstTile::DType; | ||
| 429 | + | ||
| 430 | + static_assert(IsValidGatherDType<T>::value, | ||
| 431 | + "MGATHER A5 GM->L1 data type must be int8/uint8/int16/uint16/int32/uint32/half/bfloat16/float/" | ||
| 432 | + "hifloat8/float8_e4m3/float8_e5m2."); | ||
| 433 | + static_assert(std::is_same_v<typename GlobalTable::DType, __gm__ T>, | ||
| 434 | + "MGATHER A5 GM->L1 table must be a GM GlobalTensor with element type matching the destination."); | ||
| 435 | + static_assert(std::is_same_v<typename IdxSrc::DType, __gm__ int32_t> || | ||
| 436 | + std::is_same_v<typename IdxSrc::DType, __gm__ uint32_t>, | ||
| 437 | + "MGATHER A5 GM->L1 indices must be a GM int32_t/uint32_t GlobalTensor."); | ||
| 438 | + static_assert(DstTile::Loc == TileType::Mat, "MGATHER A5 GM->L1 destination must be a Mat tile (L1)."); | ||
| 439 | + static_assert(IsMGatherNZTile<DstTile>::value, | ||
| 440 | + "MGATHER A5 GM->L1 destination must be NZ (BLayout::ColMajor + SLayout::RowMajor + " | ||
| 441 | + "SFractalSize=512)."); | ||
| 442 | + static_assert(GlobalTable::layout == Layout::ND, "MGATHER A5 GM->L1 table must use Layout::ND."); | ||
| 443 | + static_assert(DstTile::Cols % (C0_SIZE_BYTE / sizeof(T)) == 0, | ||
| 444 | + "MGATHER A5 GM->L1 destination tile Cols must be a multiple of C0 (= 32 / sizeof(T))."); | ||
| 445 | + static_assert(DstTile::Rows % FRACTAL_NZ_ROW == 0, | ||
| 446 | + "MGATHER A5 GM->L1 destination tile Rows must be a multiple of FRACTAL_NZ_ROW (16)."); | ||
| 447 | + if constexpr (Mode == Coalesce::Row) { | ||
| 448 | + static_assert(sizeof(T) <= 4, "MGATHER A5 GM->L1 Coalesce::Row supports b8/b16/b32 element types."); | ||
| 449 | + } | ||
| 450 | +} | ||
| 451 | + | ||
| 213 | template <Coalesce Mode, typename TileDst, typename GlobalTable, typename TileIdx> | 452 | template <Coalesce Mode, typename TileDst, typename GlobalTable, typename TileIdx> |
| 214 | PTO_INTERNAL void MGatherCheck(const TileDst &dst, const GlobalTable &table, const TileIdx &indices) | 453 | PTO_INTERNAL void MGatherCheck(const TileDst &dst, const GlobalTable &table, const TileIdx &indices) |
| 215 | { | 454 | { |
| @@ -272,46 +511,123 @@ template <Coalesce Mode = Coalesce::Row, GatherOOB Oob = GatherOOB::Undefined, t | |||
| 272 | PTO_INTERNAL void MGATHER_IMPL(TileDst &dst, GlobalTable &table, TileIdx &indices) | 511 | PTO_INTERNAL void MGATHER_IMPL(TileDst &dst, GlobalTable &table, TileIdx &indices) |
| 273 | { | 512 | { |
| 274 | using T = typename TileDst::DType; | 513 | using T = typename TileDst::DType; |
| 275 | - using TIdx = typename TileIdx::DType; | ||
| 276 | 514 | ||
| 277 | - MGatherCheck<Mode>(dst, table, indices); | 515 | + if constexpr (TileDst::Loc == TileType::Mat) { |
| 278 | - | 516 | + MGatherCheckGm2L1<Coalesce::Row, Oob, TileDst, GlobalTable, TileIdx>(); |
| 279 | - __gm__ const T *tablePtr = reinterpret_cast<__gm__ const T *>(table.data()); | 517 | + using TIdx = std::conditional_t<std::is_same_v<typename TileIdx::DType, __gm__ uint32_t>, uint32_t, int32_t>; |
| 280 | - | 518 | + __gm__ const T *tablePtr = reinterpret_cast<__gm__ const T *>(table.data()); |
| 281 | - constexpr int kDstValidRowS = TileDst::ValidRow; | 519 | + __gm__ const TIdx *idxPtr = reinterpret_cast<__gm__ const TIdx *>(indices.data()); |
| 282 | - constexpr int kDstValidColS = TileDst::ValidCol; | 520 | + const uint32_t validRow = dst.GetValidRow(); |
| 283 | - constexpr uint32_t kValidRowsT = (kDstValidRowS > 0) ? static_cast<uint32_t>(kDstValidRowS) : 0u; | 521 | + const uint32_t validCol = dst.GetValidCol(); |
| 284 | - constexpr uint32_t kValidColsT = (kDstValidColS > 0) ? static_cast<uint32_t>(kDstValidColS) : 0u; | 522 | + const uint32_t tableRows = |
| 285 | - | 523 | + static_cast<uint32_t>(table.GetShape(GlobalTensorDim::DIM_0) * table.GetShape(GlobalTensorDim::DIM_1) * |
| 286 | - const uint32_t validRows = dst.GetValidRow(); | 524 | + table.GetShape(GlobalTensorDim::DIM_2) * table.GetShape(GlobalTensorDim::DIM_3)); |
| 287 | - const uint32_t validCols = dst.GetValidCol(); | 525 | + const uint32_t tableRowStride = static_cast<uint32_t>(table.GetStride(GlobalTensorDim::DIM_3)); |
| 288 | - | 526 | + MGatherGm2L1RowImpl<Oob, T, TIdx, TileDst>(dst.data(), tablePtr, idxPtr, validRow, validCol, tableRows, |
| 289 | - if constexpr (Mode == Coalesce::Row) { | 527 | + tableRowStride); |
| 290 | - using TableShape = typename GlobalTable::Shape; | 528 | + return; |
| 291 | - constexpr int64_t kTableRowsS = TableShape::staticShape[3]; | ||
| 292 | - constexpr uint32_t kTableRowsT = (kTableRowsS > 0) ? static_cast<uint32_t>(kTableRowsS) : 0u; | ||
| 293 | - const uint32_t tableRows = static_cast<uint32_t>(table.GetShape(GlobalTensorDim::DIM_3)); | ||
| 294 | - MGatherRowImpl<T, TIdx, Oob, kValidRowsT, kValidColsT, kTableRowsT, TileDst, TileIdx>( | ||
| 295 | - dst.data(), tablePtr, indices.data(), validRows, validCols, tableRows); | ||
| 296 | } else { | 529 | } else { |
| 297 | - using TableShape = typename GlobalTable::Shape; | 530 | + using TIdx = typename TileIdx::DType; |
| 298 | - constexpr int64_t kTS0 = TableShape::staticShape[0]; | 531 | + |
| 299 | - constexpr int64_t kTS1 = TableShape::staticShape[1]; | 532 | + MGatherCheck<Mode>(dst, table, indices); |
| 300 | - constexpr int64_t kTS2 = TableShape::staticShape[2]; | 533 | + |
| 301 | - constexpr int64_t kTS3 = TableShape::staticShape[3]; | 534 | + __gm__ const T *tablePtr = reinterpret_cast<__gm__ const T *>(table.data()); |
| 302 | - constexpr int64_t kTS4 = TableShape::staticShape[4]; | 535 | + |
| 303 | - constexpr bool kAllStatic = (kTS0 > 0) && (kTS1 > 0) && (kTS2 > 0) && (kTS3 > 0) && (kTS4 > 0); | 536 | + constexpr int kDstValidRowS = TileDst::ValidRow; |
| 304 | - constexpr uint32_t kTableSizeT = kAllStatic ? static_cast<uint32_t>(kTS0 * kTS1 * kTS2 * kTS3 * kTS4) : 0u; | 537 | + constexpr int kDstValidColS = TileDst::ValidCol; |
| 538 | + constexpr uint32_t kValidRowsT = (kDstValidRowS > 0) ? static_cast<uint32_t>(kDstValidRowS) : 0u; | ||
| 539 | + constexpr uint32_t kValidColsT = (kDstValidColS > 0) ? static_cast<uint32_t>(kDstValidColS) : 0u; | ||
| 540 | + | ||
| 541 | + const uint32_t validRows = dst.GetValidRow(); | ||
| 542 | + const uint32_t validCols = dst.GetValidCol(); | ||
| 543 | + | ||
| 544 | + if constexpr (Mode == Coalesce::Row) { | ||
| 545 | + using TableShape = typename GlobalTable::Shape; | ||
| 546 | + constexpr int64_t kTableRowsS = TableShape::staticShape[3]; | ||
| 547 | + constexpr uint32_t kTableRowsT = (kTableRowsS > 0) ? static_cast<uint32_t>(kTableRowsS) : 0u; | ||
| 548 | + const uint32_t tableRows = static_cast<uint32_t>(table.GetShape(GlobalTensorDim::DIM_3)); | ||
| 549 | + MGatherRowImpl<T, TIdx, Oob, kValidRowsT, kValidColsT, kTableRowsT, TileDst, TileIdx>( | ||
| 550 | + dst.data(), tablePtr, indices.data(), validRows, validCols, tableRows); | ||
| 551 | + } else { | ||
| 552 | + using TableShape = typename GlobalTable::Shape; | ||
| 553 | + constexpr int64_t kTS0 = TableShape::staticShape[0]; | ||
| 554 | + constexpr int64_t kTS1 = TableShape::staticShape[1]; | ||
| 555 | + constexpr int64_t kTS2 = TableShape::staticShape[2]; | ||
| 556 | + constexpr int64_t kTS3 = TableShape::staticShape[3]; | ||
| 557 | + constexpr int64_t kTS4 = TableShape::staticShape[4]; | ||
| 558 | + constexpr bool kAllStatic = (kTS0 > 0) && (kTS1 > 0) && (kTS2 > 0) && (kTS3 > 0) && (kTS4 > 0); | ||
| 559 | + constexpr uint32_t kTableSizeT = kAllStatic ? static_cast<uint32_t>(kTS0 * kTS1 * kTS2 * kTS3 * kTS4) : 0u; | ||
| 560 | + const uint32_t tableSize = | ||
| 561 | + static_cast<uint32_t>(table.GetShape(GlobalTensorDim::DIM_0) * table.GetShape(GlobalTensorDim::DIM_1) * | ||
| 562 | + table.GetShape(GlobalTensorDim::DIM_2) * table.GetShape(GlobalTensorDim::DIM_3) * | ||
| 563 | + table.GetShape(GlobalTensorDim::DIM_4)); | ||
| 564 | + if constexpr (TileDst::ValidRow == 1 && TileDst::ValidCol == 1) { | ||
| 565 | + MGatherScalarImpl<T, TIdx, Oob, TileDst, TileIdx>(dst.data(), tablePtr, indices.data(), tableSize); | ||
| 566 | + } else { | ||
| 567 | + MGatherElemImpl<T, TIdx, Oob, kValidRowsT, kValidColsT, kTableSizeT, TileDst, TileIdx>( | ||
| 568 | + dst.data(), tablePtr, indices.data(), validRows, validCols, tableSize); | ||
| 569 | + } | ||
| 570 | + } | ||
| 571 | + } | ||
| 572 | +} | ||
| 573 | + | ||
| 574 | +template <Coalesce Mode = Coalesce::Elem, GatherOOB Oob = GatherOOB::Undefined, typename TileDst, typename GlobalTable, | ||
| 575 | + typename IdxSrc, typename GlobalScratch> | ||
| 576 | +PTO_INTERNAL void MGATHER_IMPL(TileDst &dst, GlobalTable &table, IdxSrc &indices, GlobalScratch &scratch) | ||
| 577 | +{ | ||
| 578 | + using T = typename TileDst::DType; | ||
| 579 | + | ||
| 580 | + MGatherCheckGm2L1<Coalesce::Elem, Oob, TileDst, GlobalTable, IdxSrc>(); | ||
| 581 | + static_assert(std::is_same_v<typename GlobalScratch::DType, __gm__ T>, | ||
| 582 | + "MGATHER A5 GM->L1 scratch must be a GM GlobalTensor with element type matching the destination."); | ||
| 583 | + | ||
| 584 | + using TIdx = std::conditional_t<std::is_same_v<typename IdxSrc::DType, __gm__ uint32_t>, uint32_t, int32_t>; | ||
| 585 | + __gm__ const T *tablePtr = reinterpret_cast<__gm__ const T *>(table.data()); | ||
| 586 | + __gm__ const TIdx *idxPtr = reinterpret_cast<__gm__ const TIdx *>(indices.data()); | ||
| 587 | + __gm__ T *scratchPtr = reinterpret_cast<__gm__ T *>(scratch.data()); | ||
| 588 | + | ||
| 589 | + const uint32_t validRow = dst.GetValidRow(); | ||
| 590 | + const uint32_t validCol = dst.GetValidCol(); | ||
| 591 | + const uint32_t tableSize = | ||
| 592 | + static_cast<uint32_t>(table.GetShape(GlobalTensorDim::DIM_0) * table.GetShape(GlobalTensorDim::DIM_1) * | ||
| 593 | + table.GetShape(GlobalTensorDim::DIM_2) * table.GetShape(GlobalTensorDim::DIM_3) * | ||
| 594 | + table.GetShape(GlobalTensorDim::DIM_4)); | ||
| 595 | + const uint32_t idxRowStride = static_cast<uint32_t>(indices.GetStride(GlobalTensorDim::DIM_3)); | ||
| 596 | + | ||
| 597 | + MGatherGm2L1ElemImpl<Oob, T, TIdx, TileDst>(dst.data(), tablePtr, idxPtr, scratchPtr, validRow, validCol, tableSize, | ||
| 598 | + idxRowStride); | ||
| 599 | +} | ||
| 600 | + | ||
| 601 | +template <Coalesce Mode, GatherOOB Oob, GatherExec Exec, typename TileDst, typename GlobalTable, typename IdxSrc, | ||
| 602 | + typename GlobalScratch> | ||
| 603 | +PTO_INTERNAL void MGATHER_IMPL(TileDst &dst, GlobalTable &table, IdxSrc &indices, GlobalScratch &scratch) | ||
| 604 | +{ | ||
| 605 | + using T = typename TileDst::DType; | ||
| 606 | + | ||
| 607 | + if constexpr (Exec == GatherExec::Scalar) { | ||
| 608 | + MGATHER_IMPL<Mode, Oob>(dst, table, indices, scratch); | ||
| 609 | + } else { | ||
| 610 | + static_assert(Mode == Coalesce::Elem, "MGATHER A5 GM->L1 SIMT executor is only supported for Coalesce::Elem."); | ||
| 611 | + MGatherCheckGm2L1<Coalesce::Elem, Oob, TileDst, GlobalTable, IdxSrc>(); | ||
| 612 | + static_assert(std::is_same_v<typename GlobalScratch::DType, __gm__ T>, | ||
| 613 | + "MGATHER A5 GM->L1 scratch need GM GlobalTensor with element type matching the destination"); | ||
| 614 | + | ||
| 615 | + using TIdx1 = std::conditional_t<std::is_same_v<typename IdxSrc::DType, __gm__ uint32_t>, uint32_t, int32_t>; | ||
| 616 | + __gm__ const T *tablePtr = reinterpret_cast<__gm__ const T *>(table.data()); | ||
| 617 | + __gm__ const TIdx1 *idxPtr = reinterpret_cast<__gm__ const TIdx1 *>(indices.data()); | ||
| 618 | + __gm__ T *scratchPtr = reinterpret_cast<__gm__ T *>(scratch.data()); | ||
| 619 | + | ||
| 620 | + const uint32_t validRow = dst.GetValidRow(); | ||
| 621 | + const uint32_t validCol = dst.GetValidCol(); | ||
| 305 | const uint32_t tableSize = | 622 | const uint32_t tableSize = |
| 306 | static_cast<uint32_t>(table.GetShape(GlobalTensorDim::DIM_0) * table.GetShape(GlobalTensorDim::DIM_1) * | 623 | static_cast<uint32_t>(table.GetShape(GlobalTensorDim::DIM_0) * table.GetShape(GlobalTensorDim::DIM_1) * |
| 307 | table.GetShape(GlobalTensorDim::DIM_2) * table.GetShape(GlobalTensorDim::DIM_3) * | 624 | table.GetShape(GlobalTensorDim::DIM_2) * table.GetShape(GlobalTensorDim::DIM_3) * |
| 308 | table.GetShape(GlobalTensorDim::DIM_4)); | 625 | table.GetShape(GlobalTensorDim::DIM_4)); |
| 309 | - if constexpr (TileDst::ValidRow == 1 && TileDst::ValidCol == 1) { | 626 | + const uint32_t idxRowStride = static_cast<uint32_t>(indices.GetStride(GlobalTensorDim::DIM_3)); |
| 310 | - MGatherScalarImpl<T, TIdx, Oob, TileDst, TileIdx>(dst.data(), tablePtr, indices.data(), tableSize); | 627 | + |
| 311 | - } else { | 628 | + constexpr uint8_t kSimtSyncId = 2; |
| 312 | - MGatherElemImpl<T, TIdx, Oob, kValidRowsT, kValidColsT, kTableSizeT, TileDst, TileIdx>( | 629 | + MGatherGm2L1ElemSimtImpl<Oob, T, TIdx1, TileDst, kSimtSyncId>(dst.data(), tablePtr, idxPtr, scratchPtr, |
| 313 | - dst.data(), tablePtr, indices.data(), validRows, validCols, tableSize); | 630 | + validRow, validCol, tableSize, idxRowStride); |
| 314 | - } | ||
| 315 | } | 631 | } |
| 316 | } | 632 | } |
| 317 | 633 | ||
| @@ -217,6 +217,7 @@ tconcat | |||
| 217 | textract_vec | 217 | textract_vec |
| 218 | tinsert_vec | 218 | tinsert_vec |
| 219 | mgather | 219 | mgather |
| 220 | +mgather_gm2l1 | ||
| 220 | mscatter | 221 | mscatter |
| 221 | tpushpop_cv | 222 | tpushpop_cv |
| 222 | tpushpop_vc | 223 | tpushpop_vc |
| @@ -0,0 +1,9 @@ | |||
| 1 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 2 | +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 3 | +# CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 4 | +# Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 5 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 6 | +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 7 | +# See LICENSE in the root of the software repository for the full text of the License. | ||
| 8 | + | ||
| 9 | +pto_cube_st(mgather_gm2l1) | ||
| @@ -0,0 +1,179 @@ | |||
| 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 os | ||
| 14 | +import numpy as np | ||
| 15 | + | ||
| 16 | +np.random.seed(42) | ||
| 17 | + | ||
| 18 | + | ||
| 19 | +def make_table(dtype, count): | ||
| 20 | + if np.issubdtype(dtype, np.integer): | ||
| 21 | + info = np.iinfo(dtype) | ||
| 22 | + mod = min(info.max - info.min + 1, 251) | ||
| 23 | + else: | ||
| 24 | + mod = 251 | ||
| 25 | + arr = (np.arange(1, count + 1) % mod) + 1 | ||
| 26 | + return arr.astype(dtype) | ||
| 27 | + | ||
| 28 | + | ||
| 29 | +def make_idx_random(rng, shape, max_val): | ||
| 30 | + return rng.integers(0, max_val, size=shape, dtype=np.int32) | ||
| 31 | + | ||
| 32 | + | ||
| 33 | +def make_idx_with_oob(rng, shape, table_size, oob_count): | ||
| 34 | + flat = rng.integers(0, table_size, size=int(np.prod(shape)), dtype=np.int32) | ||
| 35 | + flat[:oob_count] = rng.integers(table_size, table_size * 2, size=oob_count, dtype=np.int32) | ||
| 36 | + rng.shuffle(flat) | ||
| 37 | + return flat.reshape(shape) | ||
| 38 | + | ||
| 39 | + | ||
| 40 | +def golden_row(table, idx, dst_rows, dst_cols, oob): | ||
| 41 | + table_rows = table.shape[0] | ||
| 42 | + out = np.zeros((dst_rows, dst_cols), dtype=table.dtype) | ||
| 43 | + flat = idx.reshape(-1) | ||
| 44 | + for i in range(dst_rows): | ||
| 45 | + raw = int(flat[i]) | ||
| 46 | + if oob == "clamp": | ||
| 47 | + safe = min(max(raw, 0), table_rows - 1) | ||
| 48 | + elif oob == "wrap": | ||
| 49 | + safe = raw % table_rows | ||
| 50 | + else: | ||
| 51 | + safe = raw | ||
| 52 | + if oob == "zero" and (raw < 0 or raw >= table_rows): | ||
| 53 | + out[i, :] = 0 | ||
| 54 | + else: | ||
| 55 | + out[i, :] = table[safe, :] | ||
| 56 | + return out | ||
| 57 | + | ||
| 58 | + | ||
| 59 | +def golden_elem(table_flat, idx, oob): | ||
| 60 | + out = np.zeros_like(idx, dtype=table_flat.dtype) | ||
| 61 | + table_size = table_flat.shape[0] | ||
| 62 | + flat_idx = idx.reshape(-1) | ||
| 63 | + flat_out = out.reshape(-1) | ||
| 64 | + for i in range(flat_idx.shape[0]): | ||
| 65 | + raw = int(flat_idx[i]) | ||
| 66 | + if oob == "clamp": | ||
| 67 | + flat_out[i] = table_flat[min(max(raw, 0), table_size - 1)] | ||
| 68 | + elif oob == "wrap": | ||
| 69 | + flat_out[i] = table_flat[raw % table_size] | ||
| 70 | + elif oob == "zero": | ||
| 71 | + flat_out[i] = table_flat[raw] if 0 <= raw < table_size else 0 | ||
| 72 | + else: | ||
| 73 | + flat_out[i] = table_flat[raw] | ||
| 74 | + return flat_out.reshape(idx.shape) | ||
| 75 | + | ||
| 76 | + | ||
| 77 | +def nd_to_nz(arr_2d, c0): | ||
| 78 | + r, c = arr_2d.shape | ||
| 79 | + assert r % 16 == 0 and c % c0 == 0 | ||
| 80 | + n_block_rows = r // 16 | ||
| 81 | + n_block_cols = c // c0 | ||
| 82 | + out = np.zeros((n_block_cols, n_block_rows, 16, c0), dtype=arr_2d.dtype) | ||
| 83 | + for bc in range(n_block_cols): | ||
| 84 | + for br in range(n_block_rows): | ||
| 85 | + out[bc, br] = arr_2d[br * 16 : (br + 1) * 16, bc * c0 : (bc + 1) * c0] | ||
| 86 | + return out.reshape(-1) | ||
| 87 | + | ||
| 88 | + | ||
| 89 | +def c0_of(dtype): | ||
| 90 | + return 32 // np.dtype(dtype).itemsize | ||
| 91 | + | ||
| 92 | + | ||
| 93 | +def case_row(name, dtype, r, c, table_rows, oob="undefined", idx_kind="random"): | ||
| 94 | + rng = np.random.default_rng(hash(name) & 0xFFFFFFFF) | ||
| 95 | + table = make_table(dtype, table_rows * c).reshape(table_rows, c) | ||
| 96 | + if idx_kind == "oob": | ||
| 97 | + idx = make_idx_with_oob(rng, (r, 1), table_rows, max(1, r // 2)) | ||
| 98 | + else: | ||
| 99 | + idx = make_idx_random(rng, (r, 1), table_rows) | ||
| 100 | + golden_nd = golden_row(table, idx, r, c, oob) | ||
| 101 | + return table.reshape(-1), idx, nd_to_nz(golden_nd, c0_of(dtype)) | ||
| 102 | + | ||
| 103 | + | ||
| 104 | +def case_elem(name, dtype, r, c, table_size, oob="undefined", idx_kind="random"): | ||
| 105 | + rng = np.random.default_rng(hash(name) & 0xFFFFFFFF) | ||
| 106 | + table = make_table(dtype, table_size) | ||
| 107 | + if idx_kind == "oob": | ||
| 108 | + idx = make_idx_with_oob(rng, (r, c), table_size, max(1, (r * c) // 2)) | ||
| 109 | + else: | ||
| 110 | + idx = make_idx_random(rng, (r, c), table_size) | ||
| 111 | + golden_nd = golden_elem(table, idx, oob).reshape(r, c) | ||
| 112 | + return table, idx, nd_to_nz(golden_nd, c0_of(dtype)) | ||
| 113 | + | ||
| 114 | + | ||
| 115 | +CASES = [] | ||
| 116 | + | ||
| 117 | + | ||
| 118 | +def add(name, fn): | ||
| 119 | + CASES.append((name, fn)) | ||
| 120 | + | ||
| 121 | + | ||
| 122 | +add("MGATHERGM2L1Test.case_row_float_16x16_64rows", lambda n: case_row(n, np.float32, 16, 16, 64)) | ||
| 123 | +add("MGATHERGM2L1Test.case_row_half_16x32_64rows", lambda n: case_row(n, np.float16, 16, 32, 64)) | ||
| 124 | +add("MGATHERGM2L1Test.case_row_bfloat16_16x16_64rows", lambda n: case_row(n, np.uint16, 16, 16, 64)) | ||
| 125 | +add("MGATHERGM2L1Test.case_row_int32_16x8_32rows", lambda n: case_row(n, np.int32, 16, 8, 32)) | ||
| 126 | +add("MGATHERGM2L1Test.case_row_uint32_16x16_64rows", lambda n: case_row(n, np.uint32, 16, 16, 64)) | ||
| 127 | +add("MGATHERGM2L1Test.case_row_int16_16x16_32rows", lambda n: case_row(n, np.int16, 16, 16, 32)) | ||
| 128 | +add("MGATHERGM2L1Test.case_row_uint16_16x32_48rows", lambda n: case_row(n, np.uint16, 16, 32, 48)) | ||
| 129 | +add("MGATHERGM2L1Test.case_row_int8_16x32_64rows", lambda n: case_row(n, np.int8, 16, 32, 64)) | ||
| 130 | +add("MGATHERGM2L1Test.case_row_uint8_32x32_64rows", lambda n: case_row(n, np.uint8, 32, 32, 64)) | ||
| 131 | +add( | ||
| 132 | + "MGATHERGM2L1Test.case_row_float_clamp_16x16_8rows", | ||
| 133 | + lambda n: case_row(n, np.float32, 16, 16, 8, oob="clamp", idx_kind="oob"), | ||
| 134 | +) | ||
| 135 | +add( | ||
| 136 | + "MGATHERGM2L1Test.case_row_int32_wrap_16x8_8rows", | ||
| 137 | + lambda n: case_row(n, np.int32, 16, 8, 8, oob="wrap", idx_kind="oob"), | ||
| 138 | +) | ||
| 139 | +add( | ||
| 140 | + "MGATHERGM2L1Test.case_row_half_zero_16x16_8rows", | ||
| 141 | + lambda n: case_row(n, np.float16, 16, 16, 8, oob="zero", idx_kind="oob"), | ||
| 142 | +) | ||
| 143 | + | ||
| 144 | +add("MGATHERGM2L1Test.case_elem_float_16x16_256size", lambda n: case_elem(n, np.float32, 16, 16, 256)) | ||
| 145 | +add("MGATHERGM2L1Test.case_elem_half_16x16_256size", lambda n: case_elem(n, np.float16, 16, 16, 256)) | ||
| 146 | +add("MGATHERGM2L1Test.case_elem_bfloat16_16x16_256size", lambda n: case_elem(n, np.uint16, 16, 16, 256)) | ||
| 147 | +add("MGATHERGM2L1Test.case_elem_int32_16x8_128size", lambda n: case_elem(n, np.int32, 16, 8, 128)) | ||
| 148 | +add("MGATHERGM2L1Test.case_elem_uint32_16x16_256size", lambda n: case_elem(n, np.uint32, 16, 16, 256)) | ||
| 149 | +add("MGATHERGM2L1Test.case_elem_int16_16x16_256size", lambda n: case_elem(n, np.int16, 16, 16, 256)) | ||
| 150 | +add("MGATHERGM2L1Test.case_elem_uint16_16x32_512size", lambda n: case_elem(n, np.uint16, 16, 32, 512)) | ||
| 151 | +add("MGATHERGM2L1Test.case_elem_int8_16x32_512size", lambda n: case_elem(n, np.int8, 16, 32, 512)) | ||
| 152 | +add("MGATHERGM2L1Test.case_elem_uint8_32x32_1024size", lambda n: case_elem(n, np.uint8, 32, 32, 1024)) | ||
| 153 | +add( | ||
| 154 | + "MGATHERGM2L1Test.case_elem_float_clamp_16x16_64size", | ||
| 155 | + lambda n: case_elem(n, np.float32, 16, 16, 64, oob="clamp", idx_kind="oob"), | ||
| 156 | +) | ||
| 157 | +add( | ||
| 158 | + "MGATHERGM2L1Test.case_elem_int32_wrap_16x8_32size", | ||
| 159 | + lambda n: case_elem(n, np.int32, 16, 8, 32, oob="wrap", idx_kind="oob"), | ||
| 160 | +) | ||
| 161 | +add( | ||
| 162 | + "MGATHERGM2L1Test.case_elem_half_zero_16x16_64size", | ||
| 163 | + lambda n: case_elem(n, np.float16, 16, 16, 64, oob="zero", idx_kind="oob"), | ||
| 164 | +) | ||
| 165 | + | ||
| 166 | + | ||
| 167 | +if __name__ == "__main__": | ||
| 168 | + for name, fn in CASES: | ||
| 169 | + if not os.path.exists(name): | ||
| 170 | + os.makedirs(name) | ||
| 171 | + original_dir = os.getcwd() | ||
| 172 | + os.chdir(name) | ||
| 173 | + table, idx, golden = fn(name) | ||
| 174 | + table.tofile("table.bin") | ||
| 175 | + idx.astype(np.int32).tofile("indices.bin") | ||
| 176 | + golden.tofile("golden.bin") | ||
| 177 | + os.chdir(original_dir) | ||
| 178 | + print(f"Generated {name}") | ||
| 179 | + print("All MGATHER GM2L1 A2/A3 test data generated successfully") | ||
| @@ -0,0 +1,160 @@ | |||
| 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 | +class MGATHERGM2L1Test : public testing::Test { | ||
| 19 | +protected: | ||
| 20 | + void SetUp() override | ||
| 21 | + {} | ||
| 22 | + void TearDown() override | ||
| 23 | + {} | ||
| 24 | +}; | ||
| 25 | + | ||
| 26 | +static std::string GetGoldenDir() | ||
| 27 | +{ | ||
| 28 | + const testing::TestInfo *testInfo = testing::UnitTest::GetInstance()->current_test_info(); | ||
| 29 | + return std::string("../") + testInfo->test_suite_name() + "." + testInfo->name(); | ||
| 30 | +} | ||
| 31 | + | ||
| 32 | +template <typename T, typename TIdx, typename Launcher> | ||
| 33 | +void run_gm2l1_test(size_t tableCount, size_t idxCount, size_t outCount, Launcher launcher) | ||
| 34 | +{ | ||
| 35 | + size_t tableByteSize = tableCount * sizeof(T); | ||
| 36 | + size_t idxByteSize = idxCount * sizeof(TIdx); | ||
| 37 | + size_t outByteSize = outCount * sizeof(T); | ||
| 38 | + | ||
| 39 | + aclInit(nullptr); | ||
| 40 | + aclrtSetDevice(0); | ||
| 41 | + aclrtStream stream; | ||
| 42 | + aclrtCreateStream(&stream); | ||
| 43 | + | ||
| 44 | + T *tableHost, *outHost; | ||
| 45 | + TIdx *idxHost; | ||
| 46 | + T *tableDevice, *outDevice, *scratchDevice; | ||
| 47 | + TIdx *idxDevice; | ||
| 48 | + | ||
| 49 | + aclrtMallocHost((void **)(&tableHost), tableByteSize); | ||
| 50 | + aclrtMallocHost((void **)(&idxHost), idxByteSize); | ||
| 51 | + aclrtMallocHost((void **)(&outHost), outByteSize); | ||
| 52 | + | ||
| 53 | + aclrtMalloc((void **)&tableDevice, tableByteSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 54 | + aclrtMalloc((void **)&idxDevice, idxByteSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 55 | + aclrtMalloc((void **)&outDevice, outByteSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 56 | + aclrtMalloc((void **)&scratchDevice, outByteSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 57 | + | ||
| 58 | + ReadFile(GetGoldenDir() + "/table.bin", tableByteSize, tableHost, tableByteSize); | ||
| 59 | + ReadFile(GetGoldenDir() + "/indices.bin", idxByteSize, idxHost, idxByteSize); | ||
| 60 | + | ||
| 61 | + aclrtMemcpy(tableDevice, tableByteSize, tableHost, tableByteSize, ACL_MEMCPY_HOST_TO_DEVICE); | ||
| 62 | + aclrtMemcpy(idxDevice, idxByteSize, idxHost, idxByteSize, ACL_MEMCPY_HOST_TO_DEVICE); | ||
| 63 | + | ||
| 64 | + aclrtMemset(outDevice, outByteSize, 0, outByteSize); | ||
| 65 | + aclrtMemset(scratchDevice, outByteSize, 0, outByteSize); | ||
| 66 | + | ||
| 67 | + launcher(outDevice, tableDevice, idxDevice, scratchDevice, stream); | ||
| 68 | + | ||
| 69 | + aclrtSynchronizeStream(stream); | ||
| 70 | + aclrtMemcpy(outHost, outByteSize, outDevice, outByteSize, ACL_MEMCPY_DEVICE_TO_HOST); | ||
| 71 | + | ||
| 72 | + WriteFile(GetGoldenDir() + "/output.bin", outHost, outByteSize); | ||
| 73 | + | ||
| 74 | + aclrtFree(tableDevice); | ||
| 75 | + aclrtFree(idxDevice); | ||
| 76 | + aclrtFree(outDevice); | ||
| 77 | + aclrtFree(scratchDevice); | ||
| 78 | + | ||
| 79 | + aclrtFreeHost(tableHost); | ||
| 80 | + aclrtFreeHost(idxHost); | ||
| 81 | + aclrtFreeHost(outHost); | ||
| 82 | + aclrtDestroyStream(stream); | ||
| 83 | + aclrtResetDevice(0); | ||
| 84 | + aclFinalize(); | ||
| 85 | + | ||
| 86 | + std::vector<T> golden(outCount); | ||
| 87 | + std::vector<T> devFinal(outCount); | ||
| 88 | + ReadFile(GetGoldenDir() + "/golden.bin", outByteSize, golden.data(), outByteSize); | ||
| 89 | + ReadFile(GetGoldenDir() + "/output.bin", outByteSize, devFinal.data(), outByteSize); | ||
| 90 | + | ||
| 91 | + bool ret = ResultCmp<T>(golden, devFinal, 0.0f); | ||
| 92 | + EXPECT_TRUE(ret); | ||
| 93 | +} | ||
| 94 | + | ||
| 95 | + | ||
| 96 | + void Launch_##NAME(THOST *out, THOST *table, TIDX *indices, THOST *scratch, void *stream); | ||
| 97 | + | ||
| 98 | +DECLARE_LAUNCH(row_float_16x16_64rows, float, int32_t) | ||
| 99 | +DECLARE_LAUNCH(row_half_16x32_64rows, aclFloat16, int32_t) | ||
| 100 | +DECLARE_LAUNCH(row_bfloat16_16x16_64rows, uint16_t, int32_t) | ||
| 101 | +DECLARE_LAUNCH(row_int32_16x8_32rows, int32_t, int32_t) | ||
| 102 | +DECLARE_LAUNCH(row_uint32_16x16_64rows, uint32_t, int32_t) | ||
| 103 | +DECLARE_LAUNCH(row_int16_16x16_32rows, int16_t, int32_t) | ||
| 104 | +DECLARE_LAUNCH(row_uint16_16x32_48rows, uint16_t, int32_t) | ||
| 105 | +DECLARE_LAUNCH(row_int8_16x32_64rows, int8_t, int32_t) | ||
| 106 | +DECLARE_LAUNCH(row_uint8_32x32_64rows, uint8_t, int32_t) | ||
| 107 | +DECLARE_LAUNCH(row_float_clamp_16x16_8rows, float, int32_t) | ||
| 108 | +DECLARE_LAUNCH(row_int32_wrap_16x8_8rows, int32_t, int32_t) | ||
| 109 | +DECLARE_LAUNCH(row_half_zero_16x16_8rows, aclFloat16, int32_t) | ||
| 110 | + | ||
| 111 | +DECLARE_LAUNCH(elem_float_16x16_256size, float, int32_t) | ||
| 112 | +DECLARE_LAUNCH(elem_half_16x16_256size, aclFloat16, int32_t) | ||
| 113 | +DECLARE_LAUNCH(elem_bfloat16_16x16_256size, uint16_t, int32_t) | ||
| 114 | +DECLARE_LAUNCH(elem_int32_16x8_128size, int32_t, int32_t) | ||
| 115 | +DECLARE_LAUNCH(elem_uint32_16x16_256size, uint32_t, int32_t) | ||
| 116 | +DECLARE_LAUNCH(elem_int16_16x16_256size, int16_t, int32_t) | ||
| 117 | +DECLARE_LAUNCH(elem_uint16_16x32_512size, uint16_t, int32_t) | ||
| 118 | +DECLARE_LAUNCH(elem_int8_16x32_512size, int8_t, int32_t) | ||
| 119 | +DECLARE_LAUNCH(elem_uint8_32x32_1024size, uint8_t, int32_t) | ||
| 120 | +DECLARE_LAUNCH(elem_float_clamp_16x16_64size, float, int32_t) | ||
| 121 | +DECLARE_LAUNCH(elem_int32_wrap_16x8_32size, int32_t, int32_t) | ||
| 122 | +DECLARE_LAUNCH(elem_half_zero_16x16_64size, aclFloat16, int32_t) | ||
| 123 | + | ||
| 124 | + | ||
| 125 | + TEST_F(MGATHERGM2L1Test, case_##NAME) \ | ||
| 126 | + { \ | ||
| 127 | + run_gm2l1_test<THOST, TIDX>((size_t)TR * C, (size_t)R, (size_t)R * C, Launch_##NAME); \ | ||
| 128 | + } | ||
| 129 | + | ||
| 130 | + | ||
| 131 | + TEST_F(MGATHERGM2L1Test, case_##NAME) \ | ||
| 132 | + { \ | ||
| 133 | + run_gm2l1_test<THOST, TIDX>((size_t)TS, (size_t)R * C, (size_t)R * C, Launch_##NAME); \ | ||
| 134 | + } | ||
| 135 | + | ||
| 136 | +ROW_TEST(row_float_16x16_64rows, float, int32_t, 16, 16, 64) | ||
| 137 | +ROW_TEST(row_half_16x32_64rows, aclFloat16, int32_t, 16, 32, 64) | ||
| 138 | +ROW_TEST(row_bfloat16_16x16_64rows, uint16_t, int32_t, 16, 16, 64) | ||
| 139 | +ROW_TEST(row_int32_16x8_32rows, int32_t, int32_t, 16, 8, 32) | ||
| 140 | +ROW_TEST(row_uint32_16x16_64rows, uint32_t, int32_t, 16, 16, 64) | ||
| 141 | +ROW_TEST(row_int16_16x16_32rows, int16_t, int32_t, 16, 16, 32) | ||
| 142 | +ROW_TEST(row_uint16_16x32_48rows, uint16_t, int32_t, 16, 32, 48) | ||
| 143 | +ROW_TEST(row_int8_16x32_64rows, int8_t, int32_t, 16, 32, 64) | ||
| 144 | +ROW_TEST(row_uint8_32x32_64rows, uint8_t, int32_t, 32, 32, 64) | ||
| 145 | +ROW_TEST(row_float_clamp_16x16_8rows, float, int32_t, 16, 16, 8) | ||
| 146 | +ROW_TEST(row_int32_wrap_16x8_8rows, int32_t, int32_t, 16, 8, 8) | ||
| 147 | +ROW_TEST(row_half_zero_16x16_8rows, aclFloat16, int32_t, 16, 16, 8) | ||
| 148 | + | ||
| 149 | +ELEM_TEST(elem_float_16x16_256size, float, int32_t, 16, 16, 256) | ||
| 150 | +ELEM_TEST(elem_half_16x16_256size, aclFloat16, int32_t, 16, 16, 256) | ||
| 151 | +ELEM_TEST(elem_bfloat16_16x16_256size, uint16_t, int32_t, 16, 16, 256) | ||
| 152 | +ELEM_TEST(elem_int32_16x8_128size, int32_t, int32_t, 16, 8, 128) | ||
| 153 | +ELEM_TEST(elem_uint32_16x16_256size, uint32_t, int32_t, 16, 16, 256) | ||
| 154 | +ELEM_TEST(elem_int16_16x16_256size, int16_t, int32_t, 16, 16, 256) | ||
| 155 | +ELEM_TEST(elem_uint16_16x32_512size, uint16_t, int32_t, 16, 32, 512) | ||
| 156 | +ELEM_TEST(elem_int8_16x32_512size, int8_t, int32_t, 16, 32, 512) | ||
| 157 | +ELEM_TEST(elem_uint8_32x32_1024size, uint8_t, int32_t, 32, 32, 1024) | ||
| 158 | +ELEM_TEST(elem_float_clamp_16x16_64size, float, int32_t, 16, 16, 64) | ||
| 159 | +ELEM_TEST(elem_int32_wrap_16x8_32size, int32_t, int32_t, 16, 8, 32) | ||
| 160 | +ELEM_TEST(elem_half_zero_16x16_64size, aclFloat16, int32_t, 16, 16, 64) | ||
| @@ -0,0 +1,136 @@ | |||
| 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 | + | ||
| 16 | + | ||
| 17 | +using namespace pto; | ||
| 18 | + | ||
| 19 | +template <typename T, typename TileDataSrc> | ||
| 20 | +__tf__ PTO_INTERNAL void tf_store_nz_to_gm(__gm__ T __out__ *dst, typename TileDataSrc::TileDType __in__ src, | ||
| 21 | + uint16_t nBurst, uint16_t lenBurst, uint16_t l1Gap, uint16_t gmGap) | ||
| 22 | +{ | ||
| 23 | + copy_cbuf_to_gm(dst, __cce_get_tile_ptr(src), (uint8_t)0, nBurst, lenBurst, l1Gap, gmGap); | ||
| 24 | +} | ||
| 25 | + | ||
| 26 | +template <typename T, typename TileData> | ||
| 27 | +AICORE inline void StoreNzToGm(__gm__ T *out, TileData &src) | ||
| 28 | +{ | ||
| 29 | + constexpr uint32_t blockSizeElem = BLOCK_BYTE_SIZE / sizeof(T); | ||
| 30 | + uint32_t validRow = src.GetValidRow(); | ||
| 31 | + uint32_t validCol = src.GetValidCol(); | ||
| 32 | + uint16_t nBurst = (validCol + blockSizeElem - 1) / blockSizeElem; | ||
| 33 | + uint16_t lenBurst = validRow; | ||
| 34 | + uint16_t l1Gap = TileData::Rows - validRow; | ||
| 35 | + uint16_t gmGap = 0; | ||
| 36 | + tf_store_nz_to_gm<T, TileData>(out, src.data(), nBurst, lenBurst, l1Gap, gmGap); | ||
| 37 | +} | ||
| 38 | + | ||
| 39 | +template <pto::GatherOOB Oob, typename T, typename TIdx, uint32_t kRows, uint32_t kCols, uint32_t kTableRows> | ||
| 40 | +inline AICORE void runRowL1(__gm__ T *out, __gm__ T *table, __gm__ TIdx *indices) | ||
| 41 | +{ | ||
| 42 | + using TableShape = pto::Shape<1, 1, 1, kTableRows, kCols>; | ||
| 43 | + using TableStride = pto::Stride<1, 1, 1, kCols, 1>; | ||
| 44 | + using IdxShape = pto::Shape<1, 1, 1, 1, kRows>; | ||
| 45 | + using IdxStride = pto::Stride<1, 1, 1, kRows, 1>; | ||
| 46 | + | ||
| 47 | + GlobalTensor<T, TableShape, TableStride, Layout::ND> tableGlobal(table); | ||
| 48 | + GlobalTensor<TIdx, IdxShape, IdxStride, Layout::ND> idxGlobal(indices); | ||
| 49 | + | ||
| 50 | + using DstTile = Tile<TileType::Mat, T, kRows, kCols, BLayout::ColMajor, kRows, kCols, SLayout::RowMajor, 512>; | ||
| 51 | + DstTile dstTile; | ||
| 52 | + TASSIGN(dstTile, 0x0); | ||
| 53 | + | ||
| 54 | + MGATHER<Coalesce::Row, Oob>(dstTile, tableGlobal, idxGlobal); | ||
| 55 | + | ||
| 56 | + set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); | ||
| 57 | + wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); | ||
| 58 | + | ||
| 59 | + StoreNzToGm<T, DstTile>(out, dstTile); | ||
| 60 | +} | ||
| 61 | + | ||
| 62 | +template <pto::GatherOOB Oob, typename T, typename TIdx, uint32_t kRows, uint32_t kCols, uint32_t kTableSize> | ||
| 63 | +inline AICORE void runElemL1(__gm__ T *out, __gm__ T *table, __gm__ TIdx *indices, __gm__ T *scratch) | ||
| 64 | +{ | ||
| 65 | + using TableShape = pto::Shape<1, 1, 1, 1, kTableSize>; | ||
| 66 | + using TableStride = pto::Stride<1, 1, 1, kTableSize, 1>; | ||
| 67 | + using IdxShape = pto::Shape<1, 1, 1, kRows, kCols>; | ||
| 68 | + using IdxStride = pto::Stride<1, 1, 1, kCols, 1>; | ||
| 69 | + using ScratchShape = pto::Shape<1, 1, 1, 1, kRows * kCols>; | ||
| 70 | + using ScratchStride = pto::Stride<1, 1, 1, kRows * kCols, 1>; | ||
| 71 | + | ||
| 72 | + GlobalTensor<T, TableShape, TableStride, Layout::ND> tableGlobal(table); | ||
| 73 | + GlobalTensor<TIdx, IdxShape, IdxStride, Layout::ND> idxGlobal(indices); | ||
| 74 | + GlobalTensor<T, ScratchShape, ScratchStride, Layout::ND> scratchGlobal(scratch); | ||
| 75 | + | ||
| 76 | + using DstTile = Tile<TileType::Mat, T, kRows, kCols, BLayout::ColMajor, kRows, kCols, SLayout::RowMajor, 512>; | ||
| 77 | + DstTile dstTile; | ||
| 78 | + TASSIGN(dstTile, 0x0); | ||
| 79 | + | ||
| 80 | + MGATHER<Coalesce::Elem, Oob>(dstTile, tableGlobal, idxGlobal, scratchGlobal); | ||
| 81 | + | ||
| 82 | + set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); | ||
| 83 | + wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); | ||
| 84 | + | ||
| 85 | + StoreNzToGm<T, DstTile>(out, dstTile); | ||
| 86 | +} | ||
| 87 | + | ||
| 88 | + | ||
| 89 | + extern "C" __global__ AICORE void runMGATHER_##NAME(__gm__ T *out, __gm__ T *table, __gm__ TIDX *indices, \ | ||
| 90 | + __gm__ T *scratch) \ | ||
| 91 | + { \ | ||
| 92 | + runRowL1<pto::GatherOOB::OOB, T, TIDX, R, C, TR>(out, table, indices); \ | ||
| 93 | + } \ | ||
| 94 | + void Launch_##NAME(THOST *out, THOST *table, TIDX *indices, THOST *scratch, void *stream) \ | ||
| 95 | + { \ | ||
| 96 | + runMGATHER_##NAME<<<1, nullptr, stream>>>(reinterpret_cast<T *>(out), reinterpret_cast<T *>(table), indices, \ | ||
| 97 | + reinterpret_cast<T *>(scratch)); \ | ||
| 98 | + } | ||
| 99 | + | ||
| 100 | + | ||
| 101 | + extern "C" __global__ AICORE void runMGATHER_##NAME(__gm__ T *out, __gm__ T *table, __gm__ TIDX *indices, \ | ||
| 102 | + __gm__ T *scratch) \ | ||
| 103 | + { \ | ||
| 104 | + runElemL1<pto::GatherOOB::OOB, T, TIDX, R, C, TS>(out, table, indices, scratch); \ | ||
| 105 | + } \ | ||
| 106 | + void Launch_##NAME(THOST *out, THOST *table, TIDX *indices, THOST *scratch, void *stream) \ | ||
| 107 | + { \ | ||
| 108 | + runMGATHER_##NAME<<<1, nullptr, stream>>>(reinterpret_cast<T *>(out), reinterpret_cast<T *>(table), indices, \ | ||
| 109 | + reinterpret_cast<T *>(scratch)); \ | ||
| 110 | + } | ||
| 111 | + | ||
| 112 | +DEFINE_ROW_L1(row_float_16x16_64rows, float, float, int32_t, 16, 16, 64, Undefined) | ||
| 113 | +DEFINE_ROW_L1(row_half_16x32_64rows, aclFloat16, half, int32_t, 16, 32, 64, Undefined) | ||
| 114 | +DEFINE_ROW_L1(row_bfloat16_16x16_64rows, uint16_t, bfloat16_t, int32_t, 16, 16, 64, Undefined) | ||
| 115 | +DEFINE_ROW_L1(row_int32_16x8_32rows, int32_t, int32_t, int32_t, 16, 8, 32, Undefined) | ||
| 116 | +DEFINE_ROW_L1(row_uint32_16x16_64rows, uint32_t, uint32_t, int32_t, 16, 16, 64, Undefined) | ||
| 117 | +DEFINE_ROW_L1(row_int16_16x16_32rows, int16_t, int16_t, int32_t, 16, 16, 32, Undefined) | ||
| 118 | +DEFINE_ROW_L1(row_uint16_16x32_48rows, uint16_t, uint16_t, int32_t, 16, 32, 48, Undefined) | ||
| 119 | +DEFINE_ROW_L1(row_int8_16x32_64rows, int8_t, int8_t, int32_t, 16, 32, 64, Undefined) | ||
| 120 | +DEFINE_ROW_L1(row_uint8_32x32_64rows, uint8_t, uint8_t, int32_t, 32, 32, 64, Undefined) | ||
| 121 | +DEFINE_ROW_L1(row_float_clamp_16x16_8rows, float, float, int32_t, 16, 16, 8, Clamp) | ||
| 122 | +DEFINE_ROW_L1(row_int32_wrap_16x8_8rows, int32_t, int32_t, int32_t, 16, 8, 8, Wrap) | ||
| 123 | +DEFINE_ROW_L1(row_half_zero_16x16_8rows, aclFloat16, half, int32_t, 16, 16, 8, Zero) | ||
| 124 | + | ||
| 125 | +DEFINE_ELEM_L1(elem_float_16x16_256size, float, float, int32_t, 16, 16, 256, Undefined) | ||
| 126 | +DEFINE_ELEM_L1(elem_half_16x16_256size, aclFloat16, half, int32_t, 16, 16, 256, Undefined) | ||
| 127 | +DEFINE_ELEM_L1(elem_bfloat16_16x16_256size, uint16_t, bfloat16_t, int32_t, 16, 16, 256, Undefined) | ||
| 128 | +DEFINE_ELEM_L1(elem_int32_16x8_128size, int32_t, int32_t, int32_t, 16, 8, 128, Undefined) | ||
| 129 | +DEFINE_ELEM_L1(elem_uint32_16x16_256size, uint32_t, uint32_t, int32_t, 16, 16, 256, Undefined) | ||
| 130 | +DEFINE_ELEM_L1(elem_int16_16x16_256size, int16_t, int16_t, int32_t, 16, 16, 256, Undefined) | ||
| 131 | +DEFINE_ELEM_L1(elem_uint16_16x32_512size, uint16_t, uint16_t, int32_t, 16, 32, 512, Undefined) | ||
| 132 | +DEFINE_ELEM_L1(elem_int8_16x32_512size, int8_t, int8_t, int32_t, 16, 32, 512, Undefined) | ||
| 133 | +DEFINE_ELEM_L1(elem_uint8_32x32_1024size, uint8_t, uint8_t, int32_t, 32, 32, 1024, Undefined) | ||
| 134 | +DEFINE_ELEM_L1(elem_float_clamp_16x16_64size, float, float, int32_t, 16, 16, 64, Clamp) | ||
| 135 | +DEFINE_ELEM_L1(elem_int32_wrap_16x8_32size, int32_t, int32_t, int32_t, 16, 8, 32, Wrap) | ||
| 136 | +DEFINE_ELEM_L1(elem_half_zero_16x16_64size, aclFloat16, half, int32_t, 16, 16, 64, Zero) | ||
| @@ -224,6 +224,7 @@ set(ALL_TESTCASES | |||
| 224 | tdequant | 224 | tdequant |
| 225 | trandom | 225 | trandom |
| 226 | mgather | 226 | mgather |
| 227 | + mgather_gm2l1 | ||
| 227 | mscatter | 228 | mscatter |
| 228 | t_dhrystone | 229 | t_dhrystone |
| 229 | tpow | 230 | tpow |
| @@ -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_mix_st(mgather_gm2l1) | ||
| @@ -0,0 +1,201 @@ | |||
| 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 os | ||
| 14 | +import numpy as np | ||
| 15 | + | ||
| 16 | +np.random.seed(42) | ||
| 17 | + | ||
| 18 | + | ||
| 19 | +def make_table(dtype, count): | ||
| 20 | + if np.issubdtype(dtype, np.integer): | ||
| 21 | + info = np.iinfo(dtype) | ||
| 22 | + mod = min(info.max - info.min + 1, 251) | ||
| 23 | + else: | ||
| 24 | + mod = 251 | ||
| 25 | + arr = (np.arange(1, count + 1) % mod) + 1 | ||
| 26 | + return arr.astype(dtype) | ||
| 27 | + | ||
| 28 | + | ||
| 29 | +def make_idx_random(rng, shape, max_val): | ||
| 30 | + return rng.integers(0, max_val, size=shape, dtype=np.int32) | ||
| 31 | + | ||
| 32 | + | ||
| 33 | +def make_idx_with_oob(rng, shape, table_size, oob_count): | ||
| 34 | + flat = rng.integers(0, table_size, size=int(np.prod(shape)), dtype=np.int32) | ||
| 35 | + flat[:oob_count] = rng.integers(table_size, table_size * 2, size=oob_count, dtype=np.int32) | ||
| 36 | + rng.shuffle(flat) | ||
| 37 | + return flat.reshape(shape) | ||
| 38 | + | ||
| 39 | + | ||
| 40 | +def golden_row(table, idx, dst_rows, dst_cols, oob): | ||
| 41 | + table_rows = table.shape[0] | ||
| 42 | + out = np.zeros((dst_rows, dst_cols), dtype=table.dtype) | ||
| 43 | + flat = idx.reshape(-1) | ||
| 44 | + for i in range(dst_rows): | ||
| 45 | + raw = int(flat[i]) | ||
| 46 | + if oob == "clamp": | ||
| 47 | + safe = min(max(raw, 0), table_rows - 1) | ||
| 48 | + elif oob == "wrap": | ||
| 49 | + safe = raw % table_rows | ||
| 50 | + else: | ||
| 51 | + safe = raw | ||
| 52 | + if oob == "zero" and (raw < 0 or raw >= table_rows): | ||
| 53 | + out[i, :] = 0 | ||
| 54 | + else: | ||
| 55 | + out[i, :] = table[safe, :] | ||
| 56 | + return out | ||
| 57 | + | ||
| 58 | + | ||
| 59 | +def golden_elem(table_flat, idx, oob): | ||
| 60 | + out = np.zeros_like(idx, dtype=table_flat.dtype) | ||
| 61 | + table_size = table_flat.shape[0] | ||
| 62 | + flat_idx = idx.reshape(-1) | ||
| 63 | + flat_out = out.reshape(-1) | ||
| 64 | + for i in range(flat_idx.shape[0]): | ||
| 65 | + raw = int(flat_idx[i]) | ||
| 66 | + if oob == "clamp": | ||
| 67 | + flat_out[i] = table_flat[min(max(raw, 0), table_size - 1)] | ||
| 68 | + elif oob == "wrap": | ||
| 69 | + flat_out[i] = table_flat[raw % table_size] | ||
| 70 | + elif oob == "zero": | ||
| 71 | + flat_out[i] = table_flat[raw] if 0 <= raw < table_size else 0 | ||
| 72 | + else: | ||
| 73 | + flat_out[i] = table_flat[raw] | ||
| 74 | + return flat_out.reshape(idx.shape) | ||
| 75 | + | ||
| 76 | + | ||
| 77 | +def nd_to_nz(arr_2d, c0): | ||
| 78 | + r, c = arr_2d.shape | ||
| 79 | + assert r % 16 == 0 and c % c0 == 0 | ||
| 80 | + n_block_rows = r // 16 | ||
| 81 | + n_block_cols = c // c0 | ||
| 82 | + out = np.zeros((n_block_cols, n_block_rows, 16, c0), dtype=arr_2d.dtype) | ||
| 83 | + for bc in range(n_block_cols): | ||
| 84 | + for br in range(n_block_rows): | ||
| 85 | + out[bc, br] = arr_2d[br * 16 : (br + 1) * 16, bc * c0 : (bc + 1) * c0] | ||
| 86 | + return out.reshape(-1) | ||
| 87 | + | ||
| 88 | + | ||
| 89 | +def c0_of(dtype): | ||
| 90 | + return 32 // np.dtype(dtype).itemsize | ||
| 91 | + | ||
| 92 | + | ||
| 93 | +def case_row(name, dtype, r, c, table_rows, oob="undefined", idx_kind="random"): | ||
| 94 | + rng = np.random.default_rng(hash(name) & 0xFFFFFFFF) | ||
| 95 | + table = make_table(dtype, table_rows * c).reshape(table_rows, c) | ||
| 96 | + if idx_kind == "oob": | ||
| 97 | + idx = make_idx_with_oob(rng, (r, 1), table_rows, max(1, r // 2)) | ||
| 98 | + else: | ||
| 99 | + idx = make_idx_random(rng, (r, 1), table_rows) | ||
| 100 | + golden_nd = golden_row(table, idx, r, c, oob) | ||
| 101 | + return table.reshape(-1), idx, nd_to_nz(golden_nd, c0_of(dtype)) | ||
| 102 | + | ||
| 103 | + | ||
| 104 | +def case_elem(name, dtype, r, c, table_size, oob="undefined", idx_kind="random"): | ||
| 105 | + rng = np.random.default_rng(hash(name) & 0xFFFFFFFF) | ||
| 106 | + table = make_table(dtype, table_size) | ||
| 107 | + if idx_kind == "oob": | ||
| 108 | + idx = make_idx_with_oob(rng, (r, c), table_size, max(1, (r * c) // 2)) | ||
| 109 | + else: | ||
| 110 | + idx = make_idx_random(rng, (r, c), table_size) | ||
| 111 | + golden_nd = golden_elem(table, idx, oob).reshape(r, c) | ||
| 112 | + return table, idx, nd_to_nz(golden_nd, c0_of(dtype)) | ||
| 113 | + | ||
| 114 | + | ||
| 115 | +CASES = [] | ||
| 116 | + | ||
| 117 | + | ||
| 118 | +def add(name, fn): | ||
| 119 | + CASES.append((name, fn)) | ||
| 120 | + | ||
| 121 | + | ||
| 122 | +add("MGATHERGM2L1Test.case_row_float_16x16_64rows", lambda n: case_row(n, np.float32, 16, 16, 64)) | ||
| 123 | +add("MGATHERGM2L1Test.case_row_half_16x32_64rows", lambda n: case_row(n, np.float16, 16, 32, 64)) | ||
| 124 | +add("MGATHERGM2L1Test.case_row_bfloat16_16x16_64rows", lambda n: case_row(n, np.uint16, 16, 16, 64)) | ||
| 125 | +add("MGATHERGM2L1Test.case_row_int32_16x8_32rows", lambda n: case_row(n, np.int32, 16, 8, 32)) | ||
| 126 | +add("MGATHERGM2L1Test.case_row_uint32_16x16_64rows", lambda n: case_row(n, np.uint32, 16, 16, 64)) | ||
| 127 | +add("MGATHERGM2L1Test.case_row_int16_16x16_32rows", lambda n: case_row(n, np.int16, 16, 16, 32)) | ||
| 128 | +add("MGATHERGM2L1Test.case_row_uint16_16x32_48rows", lambda n: case_row(n, np.uint16, 16, 32, 48)) | ||
| 129 | +add("MGATHERGM2L1Test.case_row_int8_16x32_64rows", lambda n: case_row(n, np.int8, 16, 32, 64)) | ||
| 130 | +add("MGATHERGM2L1Test.case_row_uint8_32x32_64rows", lambda n: case_row(n, np.uint8, 32, 32, 64)) | ||
| 131 | +add( | ||
| 132 | + "MGATHERGM2L1Test.case_row_float_clamp_16x16_8rows", | ||
| 133 | + lambda n: case_row(n, np.float32, 16, 16, 8, oob="clamp", idx_kind="oob"), | ||
| 134 | +) | ||
| 135 | +add( | ||
| 136 | + "MGATHERGM2L1Test.case_row_int32_wrap_16x8_8rows", | ||
| 137 | + lambda n: case_row(n, np.int32, 16, 8, 8, oob="wrap", idx_kind="oob"), | ||
| 138 | +) | ||
| 139 | +add( | ||
| 140 | + "MGATHERGM2L1Test.case_row_half_zero_16x16_8rows", | ||
| 141 | + lambda n: case_row(n, np.float16, 16, 16, 8, oob="zero", idx_kind="oob"), | ||
| 142 | +) | ||
| 143 | + | ||
| 144 | +add("MGATHERGM2L1Test.case_elem_float_16x16_256size", lambda n: case_elem(n, np.float32, 16, 16, 256)) | ||
| 145 | +add("MGATHERGM2L1Test.case_elem_half_16x16_256size", lambda n: case_elem(n, np.float16, 16, 16, 256)) | ||
| 146 | +add("MGATHERGM2L1Test.case_elem_bfloat16_16x16_256size", lambda n: case_elem(n, np.uint16, 16, 16, 256)) | ||
| 147 | +add("MGATHERGM2L1Test.case_elem_int32_16x8_128size", lambda n: case_elem(n, np.int32, 16, 8, 128)) | ||
| 148 | +add("MGATHERGM2L1Test.case_elem_uint32_16x16_256size", lambda n: case_elem(n, np.uint32, 16, 16, 256)) | ||
| 149 | +add("MGATHERGM2L1Test.case_elem_int16_16x16_256size", lambda n: case_elem(n, np.int16, 16, 16, 256)) | ||
| 150 | +add("MGATHERGM2L1Test.case_elem_uint16_16x32_512size", lambda n: case_elem(n, np.uint16, 16, 32, 512)) | ||
| 151 | +add("MGATHERGM2L1Test.case_elem_int8_16x32_512size", lambda n: case_elem(n, np.int8, 16, 32, 512)) | ||
| 152 | +add("MGATHERGM2L1Test.case_elem_uint8_32x32_1024size", lambda n: case_elem(n, np.uint8, 32, 32, 1024)) | ||
| 153 | +add( | ||
| 154 | + "MGATHERGM2L1Test.case_elem_float_clamp_16x16_64size", | ||
| 155 | + lambda n: case_elem(n, np.float32, 16, 16, 64, oob="clamp", idx_kind="oob"), | ||
| 156 | +) | ||
| 157 | +add( | ||
| 158 | + "MGATHERGM2L1Test.case_elem_int32_wrap_16x8_32size", | ||
| 159 | + lambda n: case_elem(n, np.int32, 16, 8, 32, oob="wrap", idx_kind="oob"), | ||
| 160 | +) | ||
| 161 | +add( | ||
| 162 | + "MGATHERGM2L1Test.case_elem_half_zero_16x16_64size", | ||
| 163 | + lambda n: case_elem(n, np.float16, 16, 16, 64, oob="zero", idx_kind="oob"), | ||
| 164 | +) | ||
| 165 | + | ||
| 166 | +add("MGATHERGM2L1Test.case_elem_simt_float_16x16_256size", lambda n: case_elem(n, np.float32, 16, 16, 256)) | ||
| 167 | +add("MGATHERGM2L1Test.case_elem_simt_half_16x16_256size", lambda n: case_elem(n, np.float16, 16, 16, 256)) | ||
| 168 | +add("MGATHERGM2L1Test.case_elem_simt_bfloat16_16x16_256size", lambda n: case_elem(n, np.uint16, 16, 16, 256)) | ||
| 169 | +add("MGATHERGM2L1Test.case_elem_simt_int32_16x8_128size", lambda n: case_elem(n, np.int32, 16, 8, 128)) | ||
| 170 | +add("MGATHERGM2L1Test.case_elem_simt_uint32_16x16_256size", lambda n: case_elem(n, np.uint32, 16, 16, 256)) | ||
| 171 | +add("MGATHERGM2L1Test.case_elem_simt_int16_16x16_256size", lambda n: case_elem(n, np.int16, 16, 16, 256)) | ||
| 172 | +add("MGATHERGM2L1Test.case_elem_simt_uint16_16x32_512size", lambda n: case_elem(n, np.uint16, 16, 32, 512)) | ||
| 173 | +add("MGATHERGM2L1Test.case_elem_simt_int8_16x32_512size", lambda n: case_elem(n, np.int8, 16, 32, 512)) | ||
| 174 | +add("MGATHERGM2L1Test.case_elem_simt_uint8_32x32_1024size", lambda n: case_elem(n, np.uint8, 32, 32, 1024)) | ||
| 175 | +add( | ||
| 176 | + "MGATHERGM2L1Test.case_elem_simt_float_clamp_16x16_64size", | ||
| 177 | + lambda n: case_elem(n, np.float32, 16, 16, 64, oob="clamp", idx_kind="oob"), | ||
| 178 | +) | ||
| 179 | +add( | ||
| 180 | + "MGATHERGM2L1Test.case_elem_simt_int32_wrap_16x8_32size", | ||
| 181 | + lambda n: case_elem(n, np.int32, 16, 8, 32, oob="wrap", idx_kind="oob"), | ||
| 182 | +) | ||
| 183 | +add( | ||
| 184 | + "MGATHERGM2L1Test.case_elem_simt_half_zero_16x16_64size", | ||
| 185 | + lambda n: case_elem(n, np.float16, 16, 16, 64, oob="zero", idx_kind="oob"), | ||
| 186 | +) | ||
| 187 | + | ||
| 188 | + | ||
| 189 | +if __name__ == "__main__": | ||
| 190 | + for name, fn in CASES: | ||
| 191 | + if not os.path.exists(name): | ||
| 192 | + os.makedirs(name) | ||
| 193 | + original_dir = os.getcwd() | ||
| 194 | + os.chdir(name) | ||
| 195 | + table, idx, golden = fn(name) | ||
| 196 | + table.tofile("table.bin") | ||
| 197 | + idx.astype(np.int32).tofile("indices.bin") | ||
| 198 | + golden.tofile("golden.bin") | ||
| 199 | + os.chdir(original_dir) | ||
| 200 | + print(f"Generated {name}") | ||
| 201 | + print("All MGATHER GM2L1 A5 test data generated successfully") | ||
| @@ -0,0 +1,186 @@ | |||
| 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 | +class MGATHERGM2L1Test : public testing::Test { | ||
| 19 | +protected: | ||
| 20 | + void SetUp() override | ||
| 21 | + {} | ||
| 22 | + void TearDown() override | ||
| 23 | + {} | ||
| 24 | +}; | ||
| 25 | + | ||
| 26 | +static std::string GetGoldenDir() | ||
| 27 | +{ | ||
| 28 | + const testing::TestInfo *testInfo = testing::UnitTest::GetInstance()->current_test_info(); | ||
| 29 | + return std::string("../") + testInfo->test_suite_name() + "." + testInfo->name(); | ||
| 30 | +} | ||
| 31 | + | ||
| 32 | +template <typename T, typename TIdx, typename Launcher> | ||
| 33 | +void run_gm2l1_test(size_t tableCount, size_t idxCount, size_t outCount, Launcher launcher) | ||
| 34 | +{ | ||
| 35 | + size_t tableByteSize = tableCount * sizeof(T); | ||
| 36 | + size_t idxByteSize = idxCount * sizeof(TIdx); | ||
| 37 | + size_t outByteSize = outCount * sizeof(T); | ||
| 38 | + | ||
| 39 | + aclInit(nullptr); | ||
| 40 | + aclrtSetDevice(0); | ||
| 41 | + aclrtStream stream; | ||
| 42 | + aclrtCreateStream(&stream); | ||
| 43 | + | ||
| 44 | + T *tableHost, *outHost; | ||
| 45 | + TIdx *idxHost; | ||
| 46 | + T *tableDevice, *outDevice, *scratchDevice; | ||
| 47 | + TIdx *idxDevice; | ||
| 48 | + | ||
| 49 | + aclrtMallocHost((void **)(&tableHost), tableByteSize); | ||
| 50 | + aclrtMallocHost((void **)(&idxHost), idxByteSize); | ||
| 51 | + aclrtMallocHost((void **)(&outHost), outByteSize); | ||
| 52 | + | ||
| 53 | + aclrtMalloc((void **)&tableDevice, tableByteSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 54 | + aclrtMalloc((void **)&idxDevice, idxByteSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 55 | + aclrtMalloc((void **)&outDevice, outByteSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 56 | + aclrtMalloc((void **)&scratchDevice, outByteSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 57 | + | ||
| 58 | + ReadFile(GetGoldenDir() + "/table.bin", tableByteSize, tableHost, tableByteSize); | ||
| 59 | + ReadFile(GetGoldenDir() + "/indices.bin", idxByteSize, idxHost, idxByteSize); | ||
| 60 | + | ||
| 61 | + aclrtMemcpy(tableDevice, tableByteSize, tableHost, tableByteSize, ACL_MEMCPY_HOST_TO_DEVICE); | ||
| 62 | + aclrtMemcpy(idxDevice, idxByteSize, idxHost, idxByteSize, ACL_MEMCPY_HOST_TO_DEVICE); | ||
| 63 | + | ||
| 64 | + aclrtMemset(outDevice, outByteSize, 0, outByteSize); | ||
| 65 | + aclrtMemset(scratchDevice, outByteSize, 0, outByteSize); | ||
| 66 | + | ||
| 67 | + launcher(outDevice, tableDevice, idxDevice, scratchDevice, stream); | ||
| 68 | + | ||
| 69 | + aclrtSynchronizeStream(stream); | ||
| 70 | + aclrtMemcpy(outHost, outByteSize, outDevice, outByteSize, ACL_MEMCPY_DEVICE_TO_HOST); | ||
| 71 | + | ||
| 72 | + WriteFile(GetGoldenDir() + "/output.bin", outHost, outByteSize); | ||
| 73 | + | ||
| 74 | + aclrtFree(tableDevice); | ||
| 75 | + aclrtFree(idxDevice); | ||
| 76 | + aclrtFree(outDevice); | ||
| 77 | + aclrtFree(scratchDevice); | ||
| 78 | + | ||
| 79 | + aclrtFreeHost(tableHost); | ||
| 80 | + aclrtFreeHost(idxHost); | ||
| 81 | + aclrtFreeHost(outHost); | ||
| 82 | + aclrtDestroyStream(stream); | ||
| 83 | + aclrtResetDevice(0); | ||
| 84 | + aclFinalize(); | ||
| 85 | + | ||
| 86 | + std::vector<T> golden(outCount); | ||
| 87 | + std::vector<T> devFinal(outCount); | ||
| 88 | + ReadFile(GetGoldenDir() + "/golden.bin", outByteSize, golden.data(), outByteSize); | ||
| 89 | + ReadFile(GetGoldenDir() + "/output.bin", outByteSize, devFinal.data(), outByteSize); | ||
| 90 | + | ||
| 91 | + bool ret = ResultCmp<T>(golden, devFinal, 0.0f); | ||
| 92 | + EXPECT_TRUE(ret); | ||
| 93 | +} | ||
| 94 | + | ||
| 95 | + | ||
| 96 | + void Launch_##NAME(THOST *out, THOST *table, TIDX *indices, THOST *scratch, void *stream); | ||
| 97 | + | ||
| 98 | +DECLARE_LAUNCH(row_float_16x16_64rows, float, int32_t) | ||
| 99 | +DECLARE_LAUNCH(row_half_16x32_64rows, aclFloat16, int32_t) | ||
| 100 | +DECLARE_LAUNCH(row_bfloat16_16x16_64rows, uint16_t, int32_t) | ||
| 101 | +DECLARE_LAUNCH(row_int32_16x8_32rows, int32_t, int32_t) | ||
| 102 | +DECLARE_LAUNCH(row_uint32_16x16_64rows, uint32_t, int32_t) | ||
| 103 | +DECLARE_LAUNCH(row_int16_16x16_32rows, int16_t, int32_t) | ||
| 104 | +DECLARE_LAUNCH(row_uint16_16x32_48rows, uint16_t, int32_t) | ||
| 105 | +DECLARE_LAUNCH(row_int8_16x32_64rows, int8_t, int32_t) | ||
| 106 | +DECLARE_LAUNCH(row_uint8_32x32_64rows, uint8_t, int32_t) | ||
| 107 | +DECLARE_LAUNCH(row_float_clamp_16x16_8rows, float, int32_t) | ||
| 108 | +DECLARE_LAUNCH(row_int32_wrap_16x8_8rows, int32_t, int32_t) | ||
| 109 | +DECLARE_LAUNCH(row_half_zero_16x16_8rows, aclFloat16, int32_t) | ||
| 110 | + | ||
| 111 | +DECLARE_LAUNCH(elem_float_16x16_256size, float, int32_t) | ||
| 112 | +DECLARE_LAUNCH(elem_half_16x16_256size, aclFloat16, int32_t) | ||
| 113 | +DECLARE_LAUNCH(elem_bfloat16_16x16_256size, uint16_t, int32_t) | ||
| 114 | +DECLARE_LAUNCH(elem_int32_16x8_128size, int32_t, int32_t) | ||
| 115 | +DECLARE_LAUNCH(elem_uint32_16x16_256size, uint32_t, int32_t) | ||
| 116 | +DECLARE_LAUNCH(elem_int16_16x16_256size, int16_t, int32_t) | ||
| 117 | +DECLARE_LAUNCH(elem_uint16_16x32_512size, uint16_t, int32_t) | ||
| 118 | +DECLARE_LAUNCH(elem_int8_16x32_512size, int8_t, int32_t) | ||
| 119 | +DECLARE_LAUNCH(elem_uint8_32x32_1024size, uint8_t, int32_t) | ||
| 120 | +DECLARE_LAUNCH(elem_float_clamp_16x16_64size, float, int32_t) | ||
| 121 | +DECLARE_LAUNCH(elem_int32_wrap_16x8_32size, int32_t, int32_t) | ||
| 122 | +DECLARE_LAUNCH(elem_half_zero_16x16_64size, aclFloat16, int32_t) | ||
| 123 | + | ||
| 124 | +DECLARE_LAUNCH(elem_simt_float_16x16_256size, float, int32_t) | ||
| 125 | +DECLARE_LAUNCH(elem_simt_half_16x16_256size, aclFloat16, int32_t) | ||
| 126 | +DECLARE_LAUNCH(elem_simt_bfloat16_16x16_256size, uint16_t, int32_t) | ||
| 127 | +DECLARE_LAUNCH(elem_simt_int32_16x8_128size, int32_t, int32_t) | ||
| 128 | +DECLARE_LAUNCH(elem_simt_uint32_16x16_256size, uint32_t, int32_t) | ||
| 129 | +DECLARE_LAUNCH(elem_simt_int16_16x16_256size, int16_t, int32_t) | ||
| 130 | +DECLARE_LAUNCH(elem_simt_uint16_16x32_512size, uint16_t, int32_t) | ||
| 131 | +DECLARE_LAUNCH(elem_simt_int8_16x32_512size, int8_t, int32_t) | ||
| 132 | +DECLARE_LAUNCH(elem_simt_uint8_32x32_1024size, uint8_t, int32_t) | ||
| 133 | +DECLARE_LAUNCH(elem_simt_float_clamp_16x16_64size, float, int32_t) | ||
| 134 | +DECLARE_LAUNCH(elem_simt_int32_wrap_16x8_32size, int32_t, int32_t) | ||
| 135 | +DECLARE_LAUNCH(elem_simt_half_zero_16x16_64size, aclFloat16, int32_t) | ||
| 136 | + | ||
| 137 | + | ||
| 138 | + TEST_F(MGATHERGM2L1Test, case_##NAME) \ | ||
| 139 | + { \ | ||
| 140 | + run_gm2l1_test<THOST, TIDX>((size_t)TR * C, (size_t)R, (size_t)R * C, Launch_##NAME); \ | ||
| 141 | + } | ||
| 142 | + | ||
| 143 | + | ||
| 144 | + TEST_F(MGATHERGM2L1Test, case_##NAME) \ | ||
| 145 | + { \ | ||
| 146 | + run_gm2l1_test<THOST, TIDX>((size_t)TS, (size_t)R * C, (size_t)R * C, Launch_##NAME); \ | ||
| 147 | + } | ||
| 148 | + | ||
| 149 | +ROW_TEST(row_float_16x16_64rows, float, int32_t, 16, 16, 64) | ||
| 150 | +ROW_TEST(row_half_16x32_64rows, aclFloat16, int32_t, 16, 32, 64) | ||
| 151 | +ROW_TEST(row_bfloat16_16x16_64rows, uint16_t, int32_t, 16, 16, 64) | ||
| 152 | +ROW_TEST(row_int32_16x8_32rows, int32_t, int32_t, 16, 8, 32) | ||
| 153 | +ROW_TEST(row_uint32_16x16_64rows, uint32_t, int32_t, 16, 16, 64) | ||
| 154 | +ROW_TEST(row_int16_16x16_32rows, int16_t, int32_t, 16, 16, 32) | ||
| 155 | +ROW_TEST(row_uint16_16x32_48rows, uint16_t, int32_t, 16, 32, 48) | ||
| 156 | +ROW_TEST(row_int8_16x32_64rows, int8_t, int32_t, 16, 32, 64) | ||
| 157 | +ROW_TEST(row_uint8_32x32_64rows, uint8_t, int32_t, 32, 32, 64) | ||
| 158 | +ROW_TEST(row_float_clamp_16x16_8rows, float, int32_t, 16, 16, 8) | ||
| 159 | +ROW_TEST(row_int32_wrap_16x8_8rows, int32_t, int32_t, 16, 8, 8) | ||
| 160 | +ROW_TEST(row_half_zero_16x16_8rows, aclFloat16, int32_t, 16, 16, 8) | ||
| 161 | + | ||
| 162 | +ELEM_TEST(elem_float_16x16_256size, float, int32_t, 16, 16, 256) | ||
| 163 | +ELEM_TEST(elem_half_16x16_256size, aclFloat16, int32_t, 16, 16, 256) | ||
| 164 | +ELEM_TEST(elem_bfloat16_16x16_256size, uint16_t, int32_t, 16, 16, 256) | ||
| 165 | +ELEM_TEST(elem_int32_16x8_128size, int32_t, int32_t, 16, 8, 128) | ||
| 166 | +ELEM_TEST(elem_uint32_16x16_256size, uint32_t, int32_t, 16, 16, 256) | ||
| 167 | +ELEM_TEST(elem_int16_16x16_256size, int16_t, int32_t, 16, 16, 256) | ||
| 168 | +ELEM_TEST(elem_uint16_16x32_512size, uint16_t, int32_t, 16, 32, 512) | ||
| 169 | +ELEM_TEST(elem_int8_16x32_512size, int8_t, int32_t, 16, 32, 512) | ||
| 170 | +ELEM_TEST(elem_uint8_32x32_1024size, uint8_t, int32_t, 32, 32, 1024) | ||
| 171 | +ELEM_TEST(elem_float_clamp_16x16_64size, float, int32_t, 16, 16, 64) | ||
| 172 | +ELEM_TEST(elem_int32_wrap_16x8_32size, int32_t, int32_t, 16, 8, 32) | ||
| 173 | +ELEM_TEST(elem_half_zero_16x16_64size, aclFloat16, int32_t, 16, 16, 64) | ||
| 174 | + | ||
| 175 | +ELEM_TEST(elem_simt_float_16x16_256size, float, int32_t, 16, 16, 256) | ||
| 176 | +ELEM_TEST(elem_simt_half_16x16_256size, aclFloat16, int32_t, 16, 16, 256) | ||
| 177 | +ELEM_TEST(elem_simt_bfloat16_16x16_256size, uint16_t, int32_t, 16, 16, 256) | ||
| 178 | +ELEM_TEST(elem_simt_int32_16x8_128size, int32_t, int32_t, 16, 8, 128) | ||
| 179 | +ELEM_TEST(elem_simt_uint32_16x16_256size, uint32_t, int32_t, 16, 16, 256) | ||
| 180 | +ELEM_TEST(elem_simt_int16_16x16_256size, int16_t, int32_t, 16, 16, 256) | ||
| 181 | +ELEM_TEST(elem_simt_uint16_16x32_512size, uint16_t, int32_t, 16, 32, 512) | ||
| 182 | +ELEM_TEST(elem_simt_int8_16x32_512size, int8_t, int32_t, 16, 32, 512) | ||
| 183 | +ELEM_TEST(elem_simt_uint8_32x32_1024size, uint8_t, int32_t, 32, 32, 1024) | ||
| 184 | +ELEM_TEST(elem_simt_float_clamp_16x16_64size, float, int32_t, 16, 16, 64) | ||
| 185 | +ELEM_TEST(elem_simt_int32_wrap_16x8_32size, int32_t, int32_t, 16, 8, 32) | ||
| 186 | +ELEM_TEST(elem_simt_half_zero_16x16_64size, aclFloat16, int32_t, 16, 16, 64) | ||
| @@ -0,0 +1,236 @@ | |||
| 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 | + | ||
| 16 | + | ||
| 17 | +using namespace pto; | ||
| 18 | + | ||
| 19 | +template <typename TileDataDst, typename TileDataSrc> | ||
| 20 | +__tf__ PTO_INTERNAL void tf_copy_cbuf_to_ubuf(typename TileDataDst::TileDType __out__ dst, | ||
| 21 | + typename TileDataSrc::TileDType __in__ src, int vec_core, int block_count, | ||
| 22 | + int block_len, int src_stride, int dst_stride) | ||
| 23 | +{ | ||
| 24 | + copy_cbuf_to_ubuf((__ubuf__ void *)__cce_get_tile_ptr(dst), (__cbuf__ void *)__cce_get_tile_ptr(src), vec_core, | ||
| 25 | + block_count, block_len, src_stride, dst_stride); | ||
| 26 | +} | ||
| 27 | + | ||
| 28 | +template <typename DstTileData, typename SrcTileData, uint8_t syncID> | ||
| 29 | +AICORE inline void MovL1ToUbuf(DstTileData &dstTile, SrcTileData &srcTile) | ||
| 30 | +{ | ||
| 31 | + | ||
| 32 | + uint16_t blockCount = 1; | ||
| 33 | + uint16_t blockLen = DstTileData::Rows * DstTileData::Cols * sizeof(typename SrcTileData::DType) / BLOCK_BYTE_SIZE; | ||
| 34 | + | ||
| 35 | + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); | ||
| 36 | + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); | ||
| 37 | + | ||
| 38 | + tf_copy_cbuf_to_ubuf<DstTileData, SrcTileData>(dstTile.data(), srcTile.data(), 0, blockCount, blockLen, 0, 0); | ||
| 39 | + tf_copy_cbuf_to_ubuf<DstTileData, SrcTileData>(dstTile.data(), srcTile.data(), 1, blockCount, blockLen, 0, 0); | ||
| 40 | + | ||
| 41 | + set_flag(PIPE_MTE1, PIPE_MTE3, EVENT_ID0); | ||
| 42 | + wait_flag(PIPE_MTE1, PIPE_MTE3, EVENT_ID0); | ||
| 43 | + | ||
| 44 | + set_intra_block(PIPE_MTE1, syncID); | ||
| 45 | + set_intra_block(PIPE_MTE1, syncID + 16); | ||
| 46 | + | ||
| 47 | +} | ||
| 48 | + | ||
| 49 | +template <pto::GatherOOB Oob, typename T, typename TIdx, uint32_t kRows, uint32_t kCols, uint32_t kTableRows> | ||
| 50 | +inline AICORE void runRowL1(__gm__ T *out, __gm__ T *table, __gm__ TIdx *indices) | ||
| 51 | +{ | ||
| 52 | + using TableShape = pto::Shape<1, 1, 1, kTableRows, kCols>; | ||
| 53 | + using TableStride = pto::Stride<1, 1, 1, kCols, 1>; | ||
| 54 | + using IdxShape = pto::Shape<1, 1, 1, 1, kRows>; | ||
| 55 | + using IdxStride = pto::Stride<1, 1, 1, kRows, 1>; | ||
| 56 | + | ||
| 57 | + GlobalTensor<T, TableShape, TableStride, Layout::ND> tableGlobal(table); | ||
| 58 | + GlobalTensor<TIdx, IdxShape, IdxStride, Layout::ND> idxGlobal(indices); | ||
| 59 | + | ||
| 60 | + using DstTile = Tile<TileType::Mat, T, kRows, kCols, BLayout::ColMajor, kRows, kCols, SLayout::RowMajor, 512>; | ||
| 61 | + using TileUBData = Tile<TileType::Vec, T, kRows, kCols, BLayout::RowMajor, -1, -1>; | ||
| 62 | + | ||
| 63 | + DstTile dstTile; | ||
| 64 | + TASSIGN(dstTile, 0x0); | ||
| 65 | + TileUBData ubTile(kRows, kCols); | ||
| 66 | + TASSIGN(ubTile, 0x0); | ||
| 67 | + | ||
| 68 | + using GlobalDataOut = | ||
| 69 | + GlobalTensor<T, pto::Shape<1, 1, 1, kRows, kCols>, | ||
| 70 | + pto::Stride<1 * kRows * kCols, 1 * kRows * kCols, kRows * kCols, kCols, 1>, Layout::ND>; | ||
| 71 | + GlobalDataOut dstGlobal(out); | ||
| 72 | + | ||
| 73 | + MGATHER<Coalesce::Row, Oob>(dstTile, tableGlobal, idxGlobal); | ||
| 74 | + | ||
| 75 | + constexpr uint8_t syncID = 0; | ||
| 76 | + MovL1ToUbuf<TileUBData, DstTile, syncID>(ubTile, dstTile); | ||
| 77 | + | ||
| 78 | + | ||
| 79 | + wait_intra_block(PIPE_MTE3, syncID); | ||
| 80 | + TSTORE(dstGlobal, ubTile); | ||
| 81 | + | ||
| 82 | + out = dstGlobal.data(); | ||
| 83 | +} | ||
| 84 | + | ||
| 85 | +template <pto::GatherOOB Oob, typename T, typename TIdx, uint32_t kRows, uint32_t kCols, uint32_t kTableSize> | ||
| 86 | +inline AICORE void runElemL1(__gm__ T *out, __gm__ T *table, __gm__ TIdx *indices, __gm__ T *scratch) | ||
| 87 | +{ | ||
| 88 | + using TableShape = pto::Shape<1, 1, 1, 1, kTableSize>; | ||
| 89 | + using TableStride = pto::Stride<1, 1, 1, kTableSize, 1>; | ||
| 90 | + using IdxShape = pto::Shape<1, 1, 1, kRows, kCols>; | ||
| 91 | + using IdxStride = pto::Stride<1, 1, 1, kCols, 1>; | ||
| 92 | + using ScratchShape = pto::Shape<1, 1, 1, 1, kRows * kCols>; | ||
| 93 | + using ScratchStride = pto::Stride<1, 1, 1, kRows * kCols, 1>; | ||
| 94 | + | ||
| 95 | + GlobalTensor<T, TableShape, TableStride, Layout::ND> tableGlobal(table); | ||
| 96 | + GlobalTensor<TIdx, IdxShape, IdxStride, Layout::ND> idxGlobal(indices); | ||
| 97 | + GlobalTensor<T, ScratchShape, ScratchStride, Layout::ND> scratchGlobal(scratch); | ||
| 98 | + | ||
| 99 | + using DstTile = Tile<TileType::Mat, T, kRows, kCols, BLayout::ColMajor, kRows, kCols, SLayout::RowMajor, 512>; | ||
| 100 | + using TileUBData = Tile<TileType::Vec, T, kRows, kCols, BLayout::RowMajor, -1, -1>; | ||
| 101 | + | ||
| 102 | + DstTile dstTile; | ||
| 103 | + TASSIGN(dstTile, 0x0); | ||
| 104 | + TileUBData ubTile(kRows, kCols); | ||
| 105 | + TASSIGN(ubTile, 0x0); | ||
| 106 | + | ||
| 107 | + using GlobalDataOut = | ||
| 108 | + GlobalTensor<T, pto::Shape<1, 1, 1, kRows, kCols>, | ||
| 109 | + pto::Stride<1 * kRows * kCols, 1 * kRows * kCols, kRows * kCols, kCols, 1>, Layout::ND>; | ||
| 110 | + GlobalDataOut dstGlobal(out); | ||
| 111 | + | ||
| 112 | + MGATHER<Coalesce::Elem, Oob>(dstTile, tableGlobal, idxGlobal, scratchGlobal); | ||
| 113 | + | ||
| 114 | + constexpr uint8_t syncID = 0; | ||
| 115 | + MovL1ToUbuf<TileUBData, DstTile, syncID>(ubTile, dstTile); | ||
| 116 | + | ||
| 117 | + | ||
| 118 | + wait_intra_block(PIPE_MTE3, syncID); | ||
| 119 | + TSTORE(dstGlobal, ubTile); | ||
| 120 | + | ||
| 121 | + out = dstGlobal.data(); | ||
| 122 | +} | ||
| 123 | + | ||
| 124 | +template <pto::GatherOOB Oob, typename T, typename TIdx, uint32_t kRows, uint32_t kCols, uint32_t kTableSize> | ||
| 125 | +inline AICORE void runElemL1Simt(__gm__ T *out, __gm__ T *table, __gm__ TIdx *indices, __gm__ T *scratch) | ||
| 126 | +{ | ||
| 127 | + using TableShape = pto::Shape<1, 1, 1, 1, kTableSize>; | ||
| 128 | + using TableStride = pto::Stride<1, 1, 1, kTableSize, 1>; | ||
| 129 | + using IdxShape = pto::Shape<1, 1, 1, kRows, kCols>; | ||
| 130 | + using IdxStride = pto::Stride<1, 1, 1, kCols, 1>; | ||
| 131 | + using ScratchShape = pto::Shape<1, 1, 1, 1, kRows * kCols>; | ||
| 132 | + using ScratchStride = pto::Stride<1, 1, 1, kRows * kCols, 1>; | ||
| 133 | + | ||
| 134 | + GlobalTensor<T, TableShape, TableStride, Layout::ND> tableGlobal(table); | ||
| 135 | + GlobalTensor<TIdx, IdxShape, IdxStride, Layout::ND> idxGlobal(indices); | ||
| 136 | + GlobalTensor<T, ScratchShape, ScratchStride, Layout::ND> scratchGlobal(scratch); | ||
| 137 | + | ||
| 138 | + using DstTile = Tile<TileType::Mat, T, kRows, kCols, BLayout::ColMajor, kRows, kCols, SLayout::RowMajor, 512>; | ||
| 139 | + using TileUBData = Tile<TileType::Vec, T, kRows, kCols, BLayout::RowMajor, -1, -1>; | ||
| 140 | + | ||
| 141 | + DstTile dstTile; | ||
| 142 | + TASSIGN(dstTile, 0x0); | ||
| 143 | + TileUBData ubTile(kRows, kCols); | ||
| 144 | + TASSIGN(ubTile, 0x0); | ||
| 145 | + | ||
| 146 | + using GlobalDataOut = | ||
| 147 | + GlobalTensor<T, pto::Shape<1, 1, 1, kRows, kCols>, | ||
| 148 | + pto::Stride<1 * kRows * kCols, 1 * kRows * kCols, kRows * kCols, kCols, 1>, Layout::ND>; | ||
| 149 | + GlobalDataOut dstGlobal(out); | ||
| 150 | + | ||
| 151 | + MGATHER<Coalesce::Elem, Oob, GatherExec::Simt>(dstTile, tableGlobal, idxGlobal, scratchGlobal); | ||
| 152 | + | ||
| 153 | + constexpr uint8_t syncID = 0; | ||
| 154 | + MovL1ToUbuf<TileUBData, DstTile, syncID>(ubTile, dstTile); | ||
| 155 | + | ||
| 156 | + | ||
| 157 | + wait_intra_block(PIPE_MTE3, syncID); | ||
| 158 | + TSTORE(dstGlobal, ubTile); | ||
| 159 | + | ||
| 160 | + out = dstGlobal.data(); | ||
| 161 | +} | ||
| 162 | + | ||
| 163 | + | ||
| 164 | + extern "C" __global__ AICORE void runMGATHER_##NAME(__gm__ T *out, __gm__ T *table, __gm__ TIDX *indices, \ | ||
| 165 | + __gm__ T *scratch) \ | ||
| 166 | + { \ | ||
| 167 | + runRowL1<pto::GatherOOB::OOB, T, TIDX, R, C, TR>(out, table, indices); \ | ||
| 168 | + } \ | ||
| 169 | + void Launch_##NAME(THOST *out, THOST *table, TIDX *indices, THOST *scratch, void *stream) \ | ||
| 170 | + { \ | ||
| 171 | + runMGATHER_##NAME<<<1, nullptr, stream>>>(reinterpret_cast<T *>(out), reinterpret_cast<T *>(table), indices, \ | ||
| 172 | + reinterpret_cast<T *>(scratch)); \ | ||
| 173 | + } | ||
| 174 | + | ||
| 175 | + | ||
| 176 | + extern "C" __global__ AICORE void runMGATHER_##NAME(__gm__ T *out, __gm__ T *table, __gm__ TIDX *indices, \ | ||
| 177 | + __gm__ T *scratch) \ | ||
| 178 | + { \ | ||
| 179 | + runElemL1<pto::GatherOOB::OOB, T, TIDX, R, C, TS>(out, table, indices, scratch); \ | ||
| 180 | + } \ | ||
| 181 | + void Launch_##NAME(THOST *out, THOST *table, TIDX *indices, THOST *scratch, void *stream) \ | ||
| 182 | + { \ | ||
| 183 | + runMGATHER_##NAME<<<1, nullptr, stream>>>(reinterpret_cast<T *>(out), reinterpret_cast<T *>(table), indices, \ | ||
| 184 | + reinterpret_cast<T *>(scratch)); \ | ||
| 185 | + } | ||
| 186 | + | ||
| 187 | +DEFINE_ROW_L1(row_float_16x16_64rows, float, float, int32_t, 16, 16, 64, Undefined) | ||
| 188 | +DEFINE_ROW_L1(row_half_16x32_64rows, aclFloat16, half, int32_t, 16, 32, 64, Undefined) | ||
| 189 | +DEFINE_ROW_L1(row_bfloat16_16x16_64rows, uint16_t, bfloat16_t, int32_t, 16, 16, 64, Undefined) | ||
| 190 | +DEFINE_ROW_L1(row_int32_16x8_32rows, int32_t, int32_t, int32_t, 16, 8, 32, Undefined) | ||
| 191 | +DEFINE_ROW_L1(row_uint32_16x16_64rows, uint32_t, uint32_t, int32_t, 16, 16, 64, Undefined) | ||
| 192 | +DEFINE_ROW_L1(row_int16_16x16_32rows, int16_t, int16_t, int32_t, 16, 16, 32, Undefined) | ||
| 193 | +DEFINE_ROW_L1(row_uint16_16x32_48rows, uint16_t, uint16_t, int32_t, 16, 32, 48, Undefined) | ||
| 194 | +DEFINE_ROW_L1(row_int8_16x32_64rows, int8_t, int8_t, int32_t, 16, 32, 64, Undefined) | ||
| 195 | +DEFINE_ROW_L1(row_uint8_32x32_64rows, uint8_t, uint8_t, int32_t, 32, 32, 64, Undefined) | ||
| 196 | +DEFINE_ROW_L1(row_float_clamp_16x16_8rows, float, float, int32_t, 16, 16, 8, Clamp) | ||
| 197 | +DEFINE_ROW_L1(row_int32_wrap_16x8_8rows, int32_t, int32_t, int32_t, 16, 8, 8, Wrap) | ||
| 198 | +DEFINE_ROW_L1(row_half_zero_16x16_8rows, aclFloat16, half, int32_t, 16, 16, 8, Zero) | ||
| 199 | + | ||
| 200 | + | ||
| 201 | + extern "C" __global__ AICORE void runMGATHER_##NAME(__gm__ T *out, __gm__ T *table, __gm__ TIDX *indices, \ | ||
| 202 | + __gm__ T *scratch) \ | ||
| 203 | + { \ | ||
| 204 | + runElemL1Simt<pto::GatherOOB::OOB, T, TIDX, R, C, TS>(out, table, indices, scratch); \ | ||
| 205 | + } \ | ||
| 206 | + void Launch_##NAME(THOST *out, THOST *table, TIDX *indices, THOST *scratch, void *stream) \ | ||
| 207 | + { \ | ||
| 208 | + runMGATHER_##NAME<<<1, nullptr, stream>>>(reinterpret_cast<T *>(out), reinterpret_cast<T *>(table), indices, \ | ||
| 209 | + reinterpret_cast<T *>(scratch)); \ | ||
| 210 | + } | ||
| 211 | + | ||
| 212 | +DEFINE_ELEM_L1(elem_float_16x16_256size, float, float, int32_t, 16, 16, 256, Undefined) | ||
| 213 | +DEFINE_ELEM_L1(elem_half_16x16_256size, aclFloat16, half, int32_t, 16, 16, 256, Undefined) | ||
| 214 | +DEFINE_ELEM_L1(elem_bfloat16_16x16_256size, uint16_t, bfloat16_t, int32_t, 16, 16, 256, Undefined) | ||
| 215 | +DEFINE_ELEM_L1(elem_int32_16x8_128size, int32_t, int32_t, int32_t, 16, 8, 128, Undefined) | ||
| 216 | +DEFINE_ELEM_L1(elem_uint32_16x16_256size, uint32_t, uint32_t, int32_t, 16, 16, 256, Undefined) | ||
| 217 | +DEFINE_ELEM_L1(elem_int16_16x16_256size, int16_t, int16_t, int32_t, 16, 16, 256, Undefined) | ||
| 218 | +DEFINE_ELEM_L1(elem_uint16_16x32_512size, uint16_t, uint16_t, int32_t, 16, 32, 512, Undefined) | ||
| 219 | +DEFINE_ELEM_L1(elem_int8_16x32_512size, int8_t, int8_t, int32_t, 16, 32, 512, Undefined) | ||
| 220 | +DEFINE_ELEM_L1(elem_uint8_32x32_1024size, uint8_t, uint8_t, int32_t, 32, 32, 1024, Undefined) | ||
| 221 | +DEFINE_ELEM_L1(elem_float_clamp_16x16_64size, float, float, int32_t, 16, 16, 64, Clamp) | ||
| 222 | +DEFINE_ELEM_L1(elem_int32_wrap_16x8_32size, int32_t, int32_t, int32_t, 16, 8, 32, Wrap) | ||
| 223 | +DEFINE_ELEM_L1(elem_half_zero_16x16_64size, aclFloat16, half, int32_t, 16, 16, 64, Zero) | ||
| 224 | + | ||
| 225 | +DEFINE_ELEM_L1_SIMT(elem_simt_float_16x16_256size, float, float, int32_t, 16, 16, 256, Undefined) | ||
| 226 | +DEFINE_ELEM_L1_SIMT(elem_simt_half_16x16_256size, aclFloat16, half, int32_t, 16, 16, 256, Undefined) | ||
| 227 | +DEFINE_ELEM_L1_SIMT(elem_simt_bfloat16_16x16_256size, uint16_t, bfloat16_t, int32_t, 16, 16, 256, Undefined) | ||
| 228 | +DEFINE_ELEM_L1_SIMT(elem_simt_int32_16x8_128size, int32_t, int32_t, int32_t, 16, 8, 128, Undefined) | ||
| 229 | +DEFINE_ELEM_L1_SIMT(elem_simt_uint32_16x16_256size, uint32_t, uint32_t, int32_t, 16, 16, 256, Undefined) | ||
| 230 | +DEFINE_ELEM_L1_SIMT(elem_simt_int16_16x16_256size, int16_t, int16_t, int32_t, 16, 16, 256, Undefined) | ||
| 231 | +DEFINE_ELEM_L1_SIMT(elem_simt_uint16_16x32_512size, uint16_t, uint16_t, int32_t, 16, 32, 512, Undefined) | ||
| 232 | +DEFINE_ELEM_L1_SIMT(elem_simt_int8_16x32_512size, int8_t, int8_t, int32_t, 16, 32, 512, Undefined) | ||
| 233 | +DEFINE_ELEM_L1_SIMT(elem_simt_uint8_32x32_1024size, uint8_t, uint8_t, int32_t, 32, 32, 1024, Undefined) | ||
| 234 | +DEFINE_ELEM_L1_SIMT(elem_simt_float_clamp_16x16_64size, float, float, int32_t, 16, 16, 64, Clamp) | ||
| 235 | +DEFINE_ELEM_L1_SIMT(elem_simt_int32_wrap_16x8_32size, int32_t, int32_t, int32_t, 16, 8, 32, Wrap) | ||
| 236 | +DEFINE_ELEM_L1_SIMT(elem_simt_half_zero_16x16_64size, aclFloat16, half, int32_t, 16, 16, 64, Zero) | ||
| @@ -273,6 +273,11 @@ if [ "$ENABLE_A3" = "true" ]; then # A2A3 | |||
| 273 | python3 tests/script/run_st.py $ARGS -w -v a3 -t tfusedmuladdrelu -g TFUSEDMULADDRELUTest.case_float_32x128_32x192_32x256_32x127 | 273 | python3 tests/script/run_st.py $ARGS -w -v a3 -t tfusedmuladdrelu -g TFUSEDMULADDRELUTest.case_float_32x128_32x192_32x256_32x127 |
| 274 | python3 tests/script/run_st.py $ARGS -w -v a3 -t tsubrelu -g TSUBRELUTest.case_float_32x128_32x192_32x256_32x127 | 274 | python3 tests/script/run_st.py $ARGS -w -v a3 -t tsubrelu -g TSUBRELUTest.case_float_32x128_32x192_32x256_32x127 |
| 275 | python3 tests/script/run_st.py $ARGS -w -v a3 -t tmuladddst -g TMULADDDSTTest.case_float_32x128_32x192_32x256_32x127 | 275 | python3 tests/script/run_st.py $ARGS -w -v a3 -t tmuladddst -g TMULADDDSTTest.case_float_32x128_32x192_32x256_32x127 |
| 276 | + python3 tests/script/run_st.py $ARGS -w -v a3 -t mgather_gm2l1 -g MGATHERTest.case_row_float_16x16_64rows | ||
| 277 | + python3 tests/script/run_st.py $ARGS -w -v a3 -t mgather_gm2l1 -g MGATHERTest.case_row_int32_16x8_32rows | ||
| 278 | + python3 tests/script/run_st.py $ARGS -w -v a3 -t mgather_gm2l1 -g MGATHERTest.case_row_float_16x16_64rows | ||
| 279 | + python3 tests/script/run_st.py $ARGS -w -v a3 -t mgather_gm2l1 -g MGATHERTest.case_row_uint16_16x32_48rows | ||
| 280 | + python3 tests/script/run_st.py $ARGS -w -v a3 -t mgather_gm2l1 -g MGATHERTest.case_elem_bfloat16_16x16_256size | ||
| 276 | 281 | ||
| 277 | if [ "$IS_AUTO_MODE" = "false" ]; then | 282 | if [ "$IS_AUTO_MODE" = "false" ]; then |
| 278 | # this testcase has to directly call CCE intrinsics now, which won't compile for auto mode; | 283 | # this testcase has to directly call CCE intrinsics now, which won't compile for auto mode; |
| @@ -387,6 +392,7 @@ if [ "$ENABLE_A3" = "true" ]; then # A2A3 | |||
| 387 | python3 tests/script/run_st.py $ARGS -w -v a3 -t tfusedmuladdrelu | 392 | python3 tests/script/run_st.py $ARGS -w -v a3 -t tfusedmuladdrelu |
| 388 | python3 tests/script/run_st.py $ARGS -w -v a3 -t tsubrelu | 393 | python3 tests/script/run_st.py $ARGS -w -v a3 -t tsubrelu |
| 389 | python3 tests/script/run_st.py $ARGS -w -v a3 -t tmuladddst | 394 | python3 tests/script/run_st.py $ARGS -w -v a3 -t tmuladddst |
| 395 | + python3 tests/script/run_st.py $ARGS -w -v a3 -t mgather_gm2l1 | ||
| 390 | if [ "$IS_AUTO_MODE" = "false" ]; then | 396 | if [ "$IS_AUTO_MODE" = "false" ]; then |
| 391 | # this testcase has to directly call CCE intrinsics now, which won't compile for auto mode; | 397 | # this testcase has to directly call CCE intrinsics now, which won't compile for auto mode; |
| 392 | # besides, auto-sync doesn't work with CCE intrisics | 398 | # besides, auto-sync doesn't work with CCE intrisics |
| @@ -601,6 +607,15 @@ if [ "$ENABLE_A5" = "true" ]; then | |||
| 601 | python3 tests/script/run_st.py $ARGS -w -v a5 -t tfusedmuladdrelu -g TFUSEDMULADDRELUTest.case_float_32x128_32x192_32x256_32x127 | 607 | python3 tests/script/run_st.py $ARGS -w -v a5 -t tfusedmuladdrelu -g TFUSEDMULADDRELUTest.case_float_32x128_32x192_32x256_32x127 |
| 602 | python3 tests/script/run_st.py $ARGS -w -v a5 -t tsubrelu -g TSUBRELUTest.case_float_32x128_32x192_32x256_32x127 | 608 | python3 tests/script/run_st.py $ARGS -w -v a5 -t tsubrelu -g TSUBRELUTest.case_float_32x128_32x192_32x256_32x127 |
| 603 | python3 tests/script/run_st.py $ARGS -w -v a5 -t tmuladddst -g TMULADDDSTTest.case_float_32x128_32x192_32x256_32x127 | 609 | python3 tests/script/run_st.py $ARGS -w -v a5 -t tmuladddst -g TMULADDDSTTest.case_float_32x128_32x192_32x256_32x127 |
| 610 | + python3 tests/script/run_st.py $ARGS -w -v a5 -t mgather_gm2l1 -g MGATHERGM2L1Test.case_row_int32_16x8_32rows | ||
| 611 | + python3 tests/script/run_st.py $ARGS -w -v a5 -t mgather_gm2l1 -g MGATHERGM2L1Test.case_elem_int32_16x8_128size | ||
| 612 | + python3 tests/script/run_st.py $ARGS -w -v a5 -t mgather_gm2l1 -g MGATHERGM2L1Test.case_elem_simt_bfloat16_16x16_256size | ||
| 613 | + python3 tests/script/run_st.py $ARGS -w -v a5 -t mgather_gm2l1 -g MGATHERGM2L1Test.case_elem_simt_half_zero_16x16_64size | ||
| 614 | + python3 tests/script/run_st.py $ARGS -w -v a5 -t mgather_gm2l1 -g MGATHERGM2L1Test.case_row_float_16x16_64rows | ||
| 615 | + python3 tests/script/run_st.py $ARGS -w -v a5 -t mgather_gm2l1 -g MGATHERGM2L1Test.case_elem_bfloat16_16x16_256size | ||
| 616 | + python3 tests/script/run_st.py $ARGS -w -v a5 -t mgather_gm2l1 -g MGATHERGM2L1Test.case_elem_simt_int8_16x32_512size | ||
| 617 | + python3 tests/script/run_st.py $ARGS -w -v a5 -t mgather_gm2l1 -g MGATHERGM2L1Test.case_elem_simt_int16_16x16_256size | ||
| 618 | + python3 tests/script/run_st.py $ARGS -w -v a5 -t mgather_gm2l1 -g MGATHERGM2L1Test.case_elem_simt_float_clamp_16x16_64size | ||
| 604 | 619 | ||
| 605 | elif [ "$ENABLE_ALL" = "true" ]; then # 所有用例 | 620 | elif [ "$ENABLE_ALL" = "true" ]; then # 所有用例 |
| 606 | python3 tests/script/build_st.py $ARGS -v a5 -t all | 621 | python3 tests/script/build_st.py $ARGS -v a5 -t all |
| @@ -740,6 +755,7 @@ if [ "$ENABLE_A5" = "true" ]; then | |||
| 740 | python3 tests/script/run_st.py $ARGS -w -v a5 -t tfusedmuladdrelu | 755 | python3 tests/script/run_st.py $ARGS -w -v a5 -t tfusedmuladdrelu |
| 741 | python3 tests/script/run_st.py $ARGS -w -v a5 -t tsubrelu | 756 | python3 tests/script/run_st.py $ARGS -w -v a5 -t tsubrelu |
| 742 | python3 tests/script/run_st.py $ARGS -w -v a5 -t tmuladddst | 757 | python3 tests/script/run_st.py $ARGS -w -v a5 -t tmuladddst |
| 758 | + python3 tests/script/run_st.py $ARGS -w -v a5 -t mgather_gm2l1 | ||
| 743 | fi | 759 | fi |
| 744 | fi | 760 | fi |
| 745 | 761 | ||