已合并
Add Ascend950 SYRK block/kernel implementation #1194
void_ptr创建于 12 天前
Add Ascend950 SYRK block/kernel implementation #1194
已合并
共 4 个文件变更+603-0
| @@ -168,6 +168,7 @@ struct BlockPrologue { | |||
| 168 | 168 | ||
| 169 | 169 | ||
| 170 | 170 | ||
| 171 | + | ||
| 171 | 172 | ||
| 172 | 173 | ||
| 173 | 174 | ||
| @@ -0,0 +1,376 @@ | |||
| 1 | +/** | ||
| 2 | + * This program is free software, you can redistribute it and/or modify. | ||
| 3 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 4 | + * This file is a part of the CANN Open Software. | ||
| 5 | + * Licensed under CANN Open Software License Agreement Version 2.0 (the | ||
| 6 | + * "License"). Please refer to the License for details. You may not use this | ||
| 7 | + * file except in compliance with the License. THIS SOFTWARE IS PROVIDED ON AN | ||
| 8 | + * "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 9 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS | ||
| 10 | + * FOR A PARTICULAR PURPOSE. See LICENSE in the root of the software repository | ||
| 11 | + * for the full text of the License. | ||
| 12 | + */ | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + | ||
| 20 | + | ||
| 21 | + | ||
| 22 | + | ||
| 23 | + | ||
| 24 | + | ||
| 25 | + | ||
| 26 | + | ||
| 27 | + | ||
| 28 | + | ||
| 29 | + | ||
| 30 | + | ||
| 31 | + | ||
| 32 | +namespace Catlass::Gemm::Block { | ||
| 33 | + | ||
| 34 | +/** | ||
| 35 | + * @brief Block-level SYRK mmad for Ascend950: Y = X * X^T. | ||
| 36 | + * | ||
| 37 | + * Layout tags are fixed by the op contract (not user-configurable): | ||
| 38 | + * - X : RowMajor [M, K] | ||
| 39 | + * - Xt : ColumnMajor [K, M] (transpose view of the same GM buffer) | ||
| 40 | + * - Y : RowMajor [M, M] | ||
| 41 | + * | ||
| 42 | + * Dispatch is fixed to MmadPingpong<Ascend950, false> (unitFlag off; HF32 / L1-resident off). | ||
| 43 | + * Dual-write cannot use unitFlag: one mmad 0b11 pairs with only one Fixpipe 0b11. | ||
| 44 | + * L0C→GM is synchronized with M_FIX; both stores use the tile CopyL0CToGmTla wrappers | ||
| 45 | + * (RowMajor nz2nd / ColumnMajor nz2dn) with default unitFlag=0. | ||
| 46 | + * | ||
| 47 | + * Dual-write policy (see example README): | ||
| 48 | + * - diagonal block: write once (nz2nd) | ||
| 49 | + * - lower triangle: write (m,n) via nz2nd and (n,m) via nz2dn (transpose) | ||
| 50 | + */ | ||
| 51 | +template < | ||
| 52 | + class L1TileShape_, class L0TileShape_, class ElementX_, class ElementY_, | ||
| 53 | + class TileCopy_ = Gemm::Tile::PackedTileCopyTla< | ||
| 54 | + Arch::Ascend950, ElementX_, layout::RowMajor, ElementX_, layout::ColumnMajor, ElementY_, layout::RowMajor>, | ||
| 55 | + class TileMmad_ = Gemm::Tile::TileMmadTla<Arch::Ascend950, ElementX_, typename TileCopy_::LayoutTagL1A>> | ||
| 56 | +struct BlockMmadSyrkTla { | ||
| 57 | +public: | ||
| 58 | + using L1TileShape = L1TileShape_; | ||
| 59 | + using L0TileShape = L0TileShape_; | ||
| 60 | + using ElementX = ElementX_; | ||
| 61 | + using ElementXt = ElementX_; | ||
| 62 | + using ElementY = ElementY_; | ||
| 63 | + using TileCopy = TileCopy_; | ||
| 64 | + using TileMmad = TileMmad_; | ||
| 65 | + | ||
| 66 | + // Fixed by Y = X * X^T with ND RowMajor storage of X / Y. | ||
| 67 | + using LayoutTagX = layout::RowMajor; | ||
| 68 | + using LayoutTagXt = layout::ColumnMajor; | ||
| 69 | + using LayoutTagY = layout::RowMajor; | ||
| 70 | + | ||
| 71 | + static_assert( | ||
| 72 | + std::is_same_v<typename TileCopy::LayoutTagA, LayoutTagX>, | ||
| 73 | + "BlockMmadSyrkTla requires TileCopy LayoutTagA = RowMajor (X)"); | ||
| 74 | + static_assert( | ||
| 75 | + std::is_same_v<typename TileCopy::LayoutTagB, LayoutTagXt>, | ||
| 76 | + "BlockMmadSyrkTla requires TileCopy LayoutTagB = ColumnMajor (X^T)"); | ||
| 77 | + static_assert( | ||
| 78 | + std::is_same_v<typename TileCopy::LayoutTagC, LayoutTagY>, | ||
| 79 | + "BlockMmadSyrkTla requires TileCopy LayoutTagC = RowMajor (Y)"); | ||
| 80 | + | ||
| 81 | + // Dual-write requires M_FIX rather than unitFlag; HF32 / L1-resident remain default-off. | ||
| 82 | + using DispatchPolicy = Gemm::MmadPingpong<Arch::Ascend950, false>; | ||
| 83 | + using ArchTag = typename DispatchPolicy::ArchTag; | ||
| 84 | + static_assert(std::is_same_v<ArchTag, Arch::Ascend950>, "BlockMmadSyrkTla is Ascend950-only"); | ||
| 85 | + static_assert(!DispatchPolicy::ENABLE_UNIT_FLAG, "BlockMmadSyrkTla cannot use unitFlag (dual-write)"); | ||
| 86 | + static_assert(!DispatchPolicy::USE_HF32_MODE, "BlockMmadSyrkTla does not support HF32"); | ||
| 87 | + static_assert(!DispatchPolicy::ENABLE_L1_RESIDENT, "BlockMmadSyrkTla does not support L1 resident"); | ||
| 88 | + | ||
| 89 | + using LayoutX = typename TileCopy::LayoutA; | ||
| 90 | + using LayoutXt = typename TileCopy::LayoutB; | ||
| 91 | + using LayoutY = typename TileCopy::LayoutC; | ||
| 92 | + | ||
| 93 | + using ElementBias = void; | ||
| 94 | + using ElementAccumulator = | ||
| 95 | + typename Gemm::helper::ElementAccumulatorSelector<ElementX, ElementXt>::ElementAccumulator; | ||
| 96 | + | ||
| 97 | + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; | ||
| 98 | + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; | ||
| 99 | + using LayoutTagL1X = typename TileCopy::LayoutTagL1A; | ||
| 100 | + using LayoutTagL1Xt = typename TileCopy::LayoutTagL1B; | ||
| 101 | + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; | ||
| 102 | + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; | ||
| 103 | + | ||
| 104 | + static_assert( | ||
| 105 | + tla::is_tuple<L1TileShape>::value && tla::is_static<L1TileShape>::value, | ||
| 106 | + "L1TileShape must be tla::tuple and static!"); | ||
| 107 | + static_assert( | ||
| 108 | + tla::is_tuple<L0TileShape>::value && tla::is_static<L0TileShape>::value, | ||
| 109 | + "L0TileShape must be tla::tuple and static!"); | ||
| 110 | + | ||
| 111 | + static constexpr uint32_t L1X_STAGES = DispatchPolicy::L1A_STAGES; | ||
| 112 | + static constexpr uint32_t L1XT_STAGES = DispatchPolicy::L1B_STAGES; | ||
| 113 | + static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; | ||
| 114 | + static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; | ||
| 115 | + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; | ||
| 116 | + | ||
| 117 | + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); | ||
| 118 | + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); | ||
| 119 | + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); | ||
| 120 | + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); | ||
| 121 | + static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); | ||
| 122 | + static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); | ||
| 123 | + | ||
| 124 | + // L1: X tile [M, K], Xt tile [K, N] (same ElementX) | ||
| 125 | + static constexpr uint32_t L1X_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementX); | ||
| 126 | + static constexpr uint32_t L1XT_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementXt); | ||
| 127 | + // L0 tile size | ||
| 128 | + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementX); | ||
| 129 | + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementXt); | ||
| 130 | + static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator); | ||
| 131 | + | ||
| 132 | + static_assert(L0C_STAGES == 1, "BlockMmadSyrkTla uses a single L0C buffer"); | ||
| 133 | + static_assert( | ||
| 134 | + tla::detail::isRowMajor<LayoutY>::value, "BlockMmadSyrkTla requires LayoutY = RowMajor for the nz2nd store"); | ||
| 135 | + static_assert( | ||
| 136 | + L1X_TILE_SIZE * L1X_STAGES + L1XT_TILE_SIZE * L1XT_STAGES <= ArchTag::L1_SIZE, | ||
| 137 | + "L1TileShape exceeding the L1 space!"); | ||
| 138 | + static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); | ||
| 139 | + static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); | ||
| 140 | + static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); | ||
| 141 | + static_assert( | ||
| 142 | + L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, | ||
| 143 | + "BlockMmadSyrkTla requires L1 and L0 tile M/N equal (no m/n L0 loop)"); | ||
| 144 | + static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K"); | ||
| 145 | + static_assert((L1X_STAGES + L1XT_STAGES) <= 8, "L1 Buffer overflow: Exceeds the supported range of EVENT(0~7)"); | ||
| 146 | + static_assert((L0A_STAGES + L0B_STAGES) <= 8, "L0 Buffer overflow: Exceeds the supported range of EVENT_ID(0~7)"); | ||
| 147 | + | ||
| 148 | + static constexpr auto L1X_LAYOUT = | ||
| 149 | + tla::MakeLayout<ElementX, LayoutTagL1X>(tla::Int<L1_TILE_M>{}, tla::Int<L1_TILE_K>{}); | ||
| 150 | + static constexpr auto L1XT_LAYOUT = | ||
| 151 | + tla::MakeLayout<ElementXt, LayoutTagL1Xt>(tla::Int<L1_TILE_K>{}, tla::Int<L1_TILE_N>{}); | ||
| 152 | + | ||
| 153 | + CATLASS_DEVICE | ||
| 154 | + BlockMmadSyrkTla() | ||
| 155 | + {} | ||
| 156 | + | ||
| 157 | + CATLASS_DEVICE | ||
| 158 | + BlockMmadSyrkTla(Arch::Resource<ArchTag>& resource, uint32_t l1BufAddrStart = 0) | ||
| 159 | + { | ||
| 160 | + if ASCEND_IS_AIC { | ||
| 161 | + AscendC::SetHF32Mode(false); | ||
| 162 | + | ||
| 163 | + uint32_t l1XOffset = l1BufAddrStart; | ||
| 164 | + uint32_t l1XtOffset = l1BufAddrStart + L1X_TILE_SIZE * L1X_STAGES; | ||
| 165 | + | ||
| 166 | + for (uint32_t i = 0; i < L1X_STAGES; i++) { | ||
| 167 | + l1XTensorList[i] = resource.l1Buf.template GetBufferByByte<ElementX>(l1XOffset + L1X_TILE_SIZE * i); | ||
| 168 | + l1XEventList[i] = static_cast<int32_t>(i); | ||
| 169 | + AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1XEventList[i]); | ||
| 170 | + } | ||
| 171 | + for (uint32_t i = 0; i < L1XT_STAGES; i++) { | ||
| 172 | + l1XtTensorList[i] = resource.l1Buf.template GetBufferByByte<ElementXt>(l1XtOffset + L1XT_TILE_SIZE * i); | ||
| 173 | + l1XtEventList[i] = static_cast<int32_t>(i + L1X_STAGES); | ||
| 174 | + AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1XtEventList[i]); | ||
| 175 | + } | ||
| 176 | + for (uint32_t i = 0; i < L0A_STAGES; i++) { | ||
| 177 | + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte<ElementX>(L0A_TILE_SIZE * i); | ||
| 178 | + l0AEventList[i] = static_cast<int32_t>(i); | ||
| 179 | + AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0AEventList[i]); | ||
| 180 | + } | ||
| 181 | + for (uint32_t i = 0; i < L0B_STAGES; i++) { | ||
| 182 | + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte<ElementXt>(L0B_TILE_SIZE * i); | ||
| 183 | + l0BEventList[i] = static_cast<int32_t>(i + L0A_STAGES); | ||
| 184 | + AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0BEventList[i]); | ||
| 185 | + } | ||
| 186 | + for (uint32_t i = 0; i < L0C_STAGES; i++) { | ||
| 187 | + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(L0C_TILE_SIZE * i); | ||
| 188 | + l0CEventList[i] = static_cast<int32_t>(i); | ||
| 189 | + AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CEventList[i]); | ||
| 190 | + } | ||
| 191 | + } | ||
| 192 | + } | ||
| 193 | + | ||
| 194 | + CATLASS_DEVICE | ||
| 195 | + ~BlockMmadSyrkTla() | ||
| 196 | + { | ||
| 197 | + if ASCEND_IS_AIC { | ||
| 198 | + for (uint32_t i = 0; i < L1X_STAGES; i++) { | ||
| 199 | + AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1XEventList[i]); | ||
| 200 | + } | ||
| 201 | + for (uint32_t i = 0; i < L1XT_STAGES; i++) { | ||
| 202 | + AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1XtEventList[i]); | ||
| 203 | + } | ||
| 204 | + for (uint32_t i = 0; i < L0A_STAGES; i++) { | ||
| 205 | + AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0AEventList[i]); | ||
| 206 | + } | ||
| 207 | + for (uint32_t i = 0; i < L0B_STAGES; i++) { | ||
| 208 | + AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0BEventList[i]); | ||
| 209 | + } | ||
| 210 | + for (uint32_t i = 0; i < L0C_STAGES; i++) { | ||
| 211 | + AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CEventList[i]); | ||
| 212 | + } | ||
| 213 | + } | ||
| 214 | + } | ||
| 215 | + | ||
| 216 | + /** | ||
| 217 | + * @brief Block-scoped SYRK mmad + dual GM store. | ||
| 218 | + * | ||
| 219 | + * @param tensorY RowMajor tile at (m,n) — nz2nd destination | ||
| 220 | + * @param tensorYT ColumnMajor view tile at the same (m,n) coords — nz2dn lands at RowMajor (n,m) | ||
| 221 | + */ | ||
| 222 | + template <class TensorX, class TensorXt, class TensorY, class TensorYT> | ||
| 223 | + CATLASS_DEVICE void operator()( | ||
| 224 | + TensorX& tensorX, TensorXt& tensorXt, TensorY& tensorY, TensorYT& tensorYT, GemmCoord const& actualShape, | ||
| 225 | + GemmCoord const& blockCoord) | ||
| 226 | + { | ||
| 227 | + using CopyGmToL1X = typename TileCopy::template CopyGmToL1A<TensorX>; | ||
| 228 | + using CopyGmToL1Xt = typename TileCopy::template CopyGmToL1B<TensorXt>; | ||
| 229 | + CopyGmToL1X copyGmToL1X; | ||
| 230 | + CopyGmToL1Xt copyGmToL1Xt; | ||
| 231 | + | ||
| 232 | + using CopyL0CToGmNz2nd = typename TileCopy::template CopyL0CToDst<TensorY>; | ||
| 233 | + using CopyL0CToGmNz2dn = typename TileCopy::template CopyL0CToDst<TensorYT>; | ||
| 234 | + CopyL0CToGmNz2nd copyL0CToGmNz2nd; | ||
| 235 | + CopyL0CToGmNz2dn copyL0CToGmNz2dn; | ||
| 236 | + | ||
| 237 | + // L1_M/N == L0_M/N: one L0 mmad covers the whole L1 tile on M/N. | ||
| 238 | + uint32_t mActual = actualShape.m(); | ||
| 239 | + uint32_t kBlockActual = actualShape.k(); | ||
| 240 | + uint32_t nActual = actualShape.n(); | ||
| 241 | + | ||
| 242 | + auto layoutInL0C = tla::MakeLayoutL0C(mActual, nActual); | ||
| 243 | + auto tensorL0C = tla::MakeTensor(l0CTensorList[0], layoutInL0C, Arch::PositionL0C{}); | ||
| 244 | + | ||
| 245 | + uint32_t kL1Actual = min(kBlockActual, L1_TILE_K); | ||
| 246 | + // load first X tile from GM to L1 | ||
| 247 | + AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1XEventList[l1XListId]); | ||
| 248 | + auto tensorL1X = tla::MakeTensor(l1XTensorList[l1XListId], L1X_LAYOUT, Arch::PositionL1{}); | ||
| 249 | + auto tensorTileX = GetTile(tensorX, tla::MakeCoord(0, 0), tla::MakeShape(mActual, kL1Actual)); | ||
| 250 | + copyGmToL1X(tensorL1X, tensorTileX); | ||
| 251 | + AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1XEventList[l1XListId]); | ||
| 252 | + | ||
| 253 | + // load first Xt tile from GM to L1 | ||
| 254 | + AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1XtEventList[l1XtListId]); | ||
| 255 | + auto tensorL1Xt = tla::MakeTensor(l1XtTensorList[l1XtListId], L1XT_LAYOUT, Arch::PositionL1{}); | ||
| 256 | + auto tensorTileXt = GetTile(tensorXt, tla::MakeCoord(0, 0), tla::MakeShape(kL1Actual, nActual)); | ||
| 257 | + copyGmToL1Xt(tensorL1Xt, tensorTileXt); | ||
| 258 | + AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1XtEventList[l1XtListId]); | ||
| 259 | + | ||
| 260 | + // Wait until the previous L0C→GM store (or constructor prime) has released L0C. | ||
| 261 | + AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CEventList[0]); | ||
| 262 | + | ||
| 263 | + uint32_t kL1Loop = CeilDiv<L1_TILE_K>(kBlockActual); | ||
| 264 | + for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) { | ||
| 265 | + uint32_t l1XListIdNext = (l1XListId + 1 < L1X_STAGES) ? (l1XListId + 1) : 0; | ||
| 266 | + uint32_t l1XtListIdNext = (l1XtListId + 1 < L1XT_STAGES) ? (l1XtListId + 1) : 0; | ||
| 267 | + uint32_t kL1ActualNext{0}; | ||
| 268 | + if (kL1Idx < kL1Loop - 1) { | ||
| 269 | + uint32_t kL1IdxNext = kL1Idx + 1; | ||
| 270 | + kL1ActualNext = (kL1IdxNext < kL1Loop - 1) ? L1_TILE_K : (kBlockActual - kL1IdxNext * L1_TILE_K); | ||
| 271 | + | ||
| 272 | + auto tensorL1XNext = tla::MakeTensor(l1XTensorList[l1XListIdNext], L1X_LAYOUT, Arch::PositionL1{}); | ||
| 273 | + auto tensorL1XtNext = tla::MakeTensor(l1XtTensorList[l1XtListIdNext], L1XT_LAYOUT, Arch::PositionL1{}); | ||
| 274 | + auto tensorTileXNext = | ||
| 275 | + GetTile(tensorX, tla::MakeCoord(0, kL1IdxNext * L1_TILE_K), tla::MakeShape(mActual, kL1ActualNext)); | ||
| 276 | + auto tensorTileXtNext = GetTile( | ||
| 277 | + tensorXt, tla::MakeCoord(kL1IdxNext * L1_TILE_K, 0), tla::MakeShape(kL1ActualNext, nActual)); | ||
| 278 | + | ||
| 279 | + AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1XEventList[l1XListIdNext]); | ||
| 280 | + copyGmToL1X(tensorL1XNext, tensorTileXNext); | ||
| 281 | + AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1XEventList[l1XListIdNext]); | ||
| 282 | + | ||
| 283 | + AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1XtEventList[l1XtListIdNext]); | ||
| 284 | + copyGmToL1Xt(tensorL1XtNext, tensorTileXtNext); | ||
| 285 | + AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1XtEventList[l1XtListIdNext]); | ||
| 286 | + } | ||
| 287 | + | ||
| 288 | + tensorL1X = tla::MakeTensor(l1XTensorList[l1XListId], L1X_LAYOUT, Arch::PositionL1{}); | ||
| 289 | + tensorL1Xt = tla::MakeTensor(l1XtTensorList[l1XtListId], L1XT_LAYOUT, Arch::PositionL1{}); | ||
| 290 | + uint32_t kL0Loop = CeilDiv<L0_TILE_K>(kL1Actual); | ||
| 291 | + | ||
| 292 | + for (uint32_t kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) { | ||
| 293 | + uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0_TILE_K : (kL1Actual - kL0Idx * L0_TILE_K); | ||
| 294 | + | ||
| 295 | + auto layoutAInL0 = tla::MakeLayout<ElementX, LayoutTagL0A>(mActual, kL0Actual); | ||
| 296 | + auto tensorL0A = tla::MakeTensor(l0ATensorList[l0AListId], layoutAInL0, Arch::PositionL0A{}); | ||
| 297 | + auto tensorTileL1X = | ||
| 298 | + GetTile(tensorL1X, tla::MakeCoord(0, kL0Idx * L0_TILE_K), tla::MakeShape(mActual, kL0Actual)); | ||
| 299 | + | ||
| 300 | + AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0AEventList[l0AListId]); | ||
| 301 | + if (kL0Idx == 0) { | ||
| 302 | + AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1XEventList[l1XListId]); | ||
| 303 | + } | ||
| 304 | + copyL1ToL0A(tensorL0A, tensorTileL1X); | ||
| 305 | + if (kL0Idx == kL0Loop - 1) { | ||
| 306 | + AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1XEventList[l1XListId]); | ||
| 307 | + } | ||
| 308 | + | ||
| 309 | + auto layoutBInL0 = tla::MakeLayout<ElementXt, LayoutTagL0B>(kL0Actual, nActual); | ||
| 310 | + auto tensorL0B = tla::MakeTensor(l0BTensorList[l0BListId], layoutBInL0, Arch::PositionL0B{}); | ||
| 311 | + auto tensorTileL1Xt = | ||
| 312 | + GetTile(tensorL1Xt, tla::MakeCoord(kL0Idx * L0_TILE_K, 0), tla::MakeShape(kL0Actual, nActual)); | ||
| 313 | + | ||
| 314 | + AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0BEventList[l0BListId]); | ||
| 315 | + if (kL0Idx == 0) { | ||
| 316 | + AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1XtEventList[l1XtListId]); | ||
| 317 | + } | ||
| 318 | + copyL1ToL0B(tensorL0B, tensorTileL1Xt); | ||
| 319 | + if (kL0Idx == kL0Loop - 1) { | ||
| 320 | + AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1XtEventList[l1XtListId]); | ||
| 321 | + } | ||
| 322 | + | ||
| 323 | + AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(l0CEventList[0]); | ||
| 324 | + AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(l0CEventList[0]); | ||
| 325 | + | ||
| 326 | + bool initC = ((kL1Idx == 0) && (kL0Idx == 0)); | ||
| 327 | + tileMmad(tensorL0C, tensorL0A, tensorL0B, mActual, nActual, kL0Actual, initC); | ||
| 328 | + | ||
| 329 | + AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0BEventList[l0BListId]); | ||
| 330 | + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; | ||
| 331 | + AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0AEventList[l0AListId]); | ||
| 332 | + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; | ||
| 333 | + } | ||
| 334 | + l1XListId = l1XListIdNext; | ||
| 335 | + l1XtListId = l1XtListIdNext; | ||
| 336 | + kL1Actual = kL1ActualNext; | ||
| 337 | + } | ||
| 338 | + | ||
| 339 | + // Dual-write: one M_FIX covers both stores; unitFlag stays 0 (tile-wrapper default). | ||
| 340 | + AscendC::SetFlag<AscendC::HardEvent::M_FIX>(l0CEventList[0]); | ||
| 341 | + AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(l0CEventList[0]); | ||
| 342 | + if (blockCoord.m() == blockCoord.n()) { | ||
| 343 | + copyL0CToGmNz2nd(tensorY, tensorL0C); | ||
| 344 | + } else { | ||
| 345 | + copyL0CToGmNz2nd(tensorY, tensorL0C); | ||
| 346 | + copyL0CToGmNz2dn(tensorYT, tensorL0C); | ||
| 347 | + } | ||
| 348 | + AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CEventList[0]); | ||
| 349 | + } | ||
| 350 | + | ||
| 351 | +protected: | ||
| 352 | + AscendC::LocalTensor<ElementX> l1XTensorList[L1X_STAGES]; | ||
| 353 | + AscendC::LocalTensor<ElementXt> l1XtTensorList[L1XT_STAGES]; | ||
| 354 | + AscendC::LocalTensor<ElementX> l0ATensorList[L0A_STAGES]; | ||
| 355 | + AscendC::LocalTensor<ElementXt> l0BTensorList[L0B_STAGES]; | ||
| 356 | + AscendC::LocalTensor<ElementAccumulator> l0CTensorList[L0C_STAGES]; | ||
| 357 | + | ||
| 358 | + int32_t l1XEventList[L1X_STAGES]; | ||
| 359 | + int32_t l1XtEventList[L1XT_STAGES]; | ||
| 360 | + int32_t l0AEventList[L0A_STAGES]; | ||
| 361 | + int32_t l0BEventList[L0B_STAGES]; | ||
| 362 | + int32_t l0CEventList[L0C_STAGES]; | ||
| 363 | + | ||
| 364 | + uint32_t l1XListId{0}; | ||
| 365 | + uint32_t l1XtListId{0}; | ||
| 366 | + uint32_t l0AListId{0}; | ||
| 367 | + uint32_t l0BListId{0}; | ||
| 368 | + | ||
| 369 | + TileMmad tileMmad; | ||
| 370 | + CopyL1ToL0A copyL1ToL0A; | ||
| 371 | + CopyL1ToL0B copyL1ToL0B; | ||
| 372 | +}; | ||
| 373 | + | ||
| 374 | +} // namespace Catlass::Gemm::Block | ||
| 375 | + | ||
| 376 | + | ||
| @@ -0,0 +1,183 @@ | |||
| 1 | +/** | ||
| 2 | + * This program is free software, you can redistribute it and/or modify. | ||
| 3 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 4 | + * This file is a part of the CANN Open Software. | ||
| 5 | + * Licensed under CANN Open Software License Agreement Version 2.0 (the | ||
| 6 | + * "License"). Please refer to the License for details. You may not use this | ||
| 7 | + * file except in compliance with the License. THIS SOFTWARE IS PROVIDED ON AN | ||
| 8 | + * "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 9 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS | ||
| 10 | + * FOR A PARTICULAR PURPOSE. See LICENSE in the root of the software repository | ||
| 11 | + * for the full text of the License. | ||
| 12 | + */ | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + | ||
| 20 | + | ||
| 21 | + | ||
| 22 | + | ||
| 23 | + | ||
| 24 | + | ||
| 25 | + | ||
| 26 | + | ||
| 27 | + | ||
| 28 | + | ||
| 29 | +namespace Catlass::Gemm::Kernel { | ||
| 30 | + | ||
| 31 | +/** | ||
| 32 | + * @brief Kernel-level SYRK: Y = X * X^T (Ascend950 only). | ||
| 33 | + * | ||
| 34 | + * Scheduling rules (lower-triangle only): | ||
| 35 | + * 1. blockCoord.m() < blockCoord.n() -> skip | ||
| 36 | + * 2. blockCoord.m() == blockCoord.n() -> compute & write once | ||
| 37 | + * 3. blockCoord.m() > blockCoord.n() -> compute & dual-write (nz2nd + nz2dn) | ||
| 38 | + * | ||
| 39 | + */ | ||
| 40 | +template <class BlockMmad_, class BlockEpilogue_, class BlockScheduler_> | ||
| 41 | +class BasicSyrkTla { | ||
| 42 | +public: | ||
| 43 | + using BlockMmad = BlockMmad_; | ||
| 44 | + using BlockEpilogue = BlockEpilogue_; | ||
| 45 | + using BlockScheduler = BlockScheduler_; | ||
| 46 | + | ||
| 47 | + using ArchTag = typename BlockMmad::ArchTag; | ||
| 48 | + using L1TileShape = typename BlockMmad::L1TileShape; | ||
| 49 | + using ElementX = typename BlockMmad::ElementX; | ||
| 50 | + using LayoutTagX = typename BlockMmad::LayoutTagX; | ||
| 51 | + using LayoutX = typename BlockMmad::LayoutX; | ||
| 52 | + using ElementXt = typename BlockMmad::ElementXt; | ||
| 53 | + using LayoutTagXt = typename BlockMmad::LayoutTagXt; | ||
| 54 | + using LayoutXt = typename BlockMmad::LayoutXt; | ||
| 55 | + using ElementY = typename BlockMmad::ElementY; | ||
| 56 | + using LayoutTagY = typename BlockMmad::LayoutTagY; | ||
| 57 | + using LayoutY = typename BlockMmad::LayoutY; | ||
| 58 | + using LayoutTagYT = layout::ColumnMajor; // ColumnMajor view of the same Y buffer for nz2dn (transpose) store. | ||
| 59 | + using LayoutYT = detail::TagToLayout_t<ElementY, LayoutTagYT>; | ||
| 60 | + using ElementAccumulator = typename BlockMmad::ElementAccumulator; | ||
| 61 | + | ||
| 62 | + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); | ||
| 63 | + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); | ||
| 64 | + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); | ||
| 65 | + | ||
| 66 | + struct Params { | ||
| 67 | + GemmCoord problemShape; // (M, M, K) | ||
| 68 | + GM_ADDR ptrX; | ||
| 69 | + GM_ADDR ptrY; | ||
| 70 | + LayoutX layoutX; | ||
| 71 | + LayoutXt layoutXt; | ||
| 72 | + LayoutY layoutY; | ||
| 73 | + LayoutYT layoutYT; | ||
| 74 | + | ||
| 75 | + CATLASS_HOST_DEVICE | ||
| 76 | + Params() | ||
| 77 | + {} | ||
| 78 | + | ||
| 79 | + CATLASS_HOST_DEVICE | ||
| 80 | + Params( | ||
| 81 | + GemmCoord const& problemShape_, GM_ADDR ptrX_, GM_ADDR ptrY_, LayoutX layoutX_, LayoutXt layoutXt_, | ||
| 82 | + LayoutY layoutY_, LayoutYT layoutYT_) | ||
| 83 | + : problemShape(problemShape_), | ||
| 84 | + ptrX(ptrX_), | ||
| 85 | + ptrY(ptrY_), | ||
| 86 | + layoutX(layoutX_), | ||
| 87 | + layoutXt(layoutXt_), | ||
| 88 | + layoutY(layoutY_), | ||
| 89 | + layoutYT(layoutYT_) | ||
| 90 | + {} | ||
| 91 | + }; | ||
| 92 | + | ||
| 93 | + struct Arguments { | ||
| 94 | + GemmCoord problemShape; // (M, M, K) | ||
| 95 | + uint8_t* ptrX; | ||
| 96 | + uint8_t* ptrY; | ||
| 97 | + }; | ||
| 98 | + | ||
| 99 | + static bool CanImplement(const Arguments& args) | ||
| 100 | + { | ||
| 101 | + return args.problemShape.m() == args.problemShape.n(); | ||
| 102 | + } | ||
| 103 | + | ||
| 104 | + static size_t GetWorkspaceSize(const Arguments& /*args*/) | ||
| 105 | + { | ||
| 106 | + return 0; | ||
| 107 | + } | ||
| 108 | + | ||
| 109 | + static Params ToUnderlyingArguments(const Arguments& args, uint8_t* /*workspace*/) | ||
| 110 | + { | ||
| 111 | + uint32_t m = args.problemShape.m(); | ||
| 112 | + uint32_t k = args.problemShape.k(); | ||
| 113 | + return Params{ | ||
| 114 | + args.problemShape, | ||
| 115 | + args.ptrX, | ||
| 116 | + args.ptrY, | ||
| 117 | + tla::MakeLayout<ElementX, LayoutTagX>(m, k), | ||
| 118 | + tla::MakeLayout<ElementXt, LayoutTagXt>(k, m), | ||
| 119 | + tla::MakeLayout<ElementY, LayoutTagY>(m, m), | ||
| 120 | + tla::MakeLayout<ElementY, LayoutTagYT>(m, m), | ||
| 121 | + }; | ||
| 122 | + } | ||
| 123 | + | ||
| 124 | + CATLASS_DEVICE | ||
| 125 | + BasicSyrkTla() | ||
| 126 | + {} | ||
| 127 | + | ||
| 128 | + template <int32_t CoreType_ = g_coreType> | ||
| 129 | + CATLASS_DEVICE void operator()(Params const& params); | ||
| 130 | + | ||
| 131 | + template <> | ||
| 132 | + CATLASS_DEVICE void operator()<AscendC::AIC>(Params const& params) | ||
| 133 | + { | ||
| 134 | + BlockScheduler blockScheduler(params.problemShape, MakeCoord(L1_TILE_M, L1_TILE_N)); | ||
| 135 | + uint32_t coreLoops = blockScheduler.GetCoreLoops(); | ||
| 136 | + | ||
| 137 | + Arch::Resource<ArchTag> resource; | ||
| 138 | + BlockMmad blockMmad(resource); | ||
| 139 | + | ||
| 140 | + AscendC::GlobalTensor<ElementX> gmX; | ||
| 141 | + gmX.SetGlobalBuffer((__gm__ ElementX*)params.ptrX); | ||
| 142 | + AscendC::GlobalTensor<ElementY> gmY; | ||
| 143 | + gmY.SetGlobalBuffer((__gm__ ElementY*)params.ptrY); | ||
| 144 | + | ||
| 145 | + auto tensorX = tla::MakeTensor(gmX, params.layoutX, Arch::PositionGM{}); | ||
| 146 | + auto tensorXt = tla::MakeTensor(gmX, params.layoutXt, Arch::PositionGM{}); | ||
| 147 | + auto tensorY = tla::MakeTensor(gmY, params.layoutY, Arch::PositionGM{}); | ||
| 148 | + auto tensorYT = tla::MakeTensor(gmY, params.layoutYT, Arch::PositionGM{}); | ||
| 149 | + | ||
| 150 | + for (uint32_t loopIdx = AscendC::GetBlockIdx(); loopIdx < coreLoops; loopIdx += AscendC::GetBlockNum()) { | ||
| 151 | + GemmCoord blockCoord = blockScheduler.GetBlockCoord(loopIdx); | ||
| 152 | + if (blockCoord.m() < blockCoord.n()) { | ||
| 153 | + continue; | ||
| 154 | + } | ||
| 155 | + | ||
| 156 | + GemmCoord actualBlockShape = blockScheduler.GetActualBlockShape(blockCoord); | ||
| 157 | + auto tileX = tla::GetTile( | ||
| 158 | + tensorX, tla::MakeCoord(blockCoord.m() * L1_TILE_M, blockCoord.k() * L1_TILE_K), | ||
| 159 | + tla::MakeShape(actualBlockShape.m(), actualBlockShape.k())); | ||
| 160 | + auto tileXt = tla::GetTile( | ||
| 161 | + tensorXt, tla::MakeCoord(blockCoord.k() * L1_TILE_K, blockCoord.n() * L1_TILE_N), | ||
| 162 | + tla::MakeShape(actualBlockShape.k(), actualBlockShape.n())); | ||
| 163 | + auto tileY = tla::GetTile( | ||
| 164 | + tensorY, tla::MakeCoord(blockCoord.m() * L1_TILE_M, blockCoord.n() * L1_TILE_N), | ||
| 165 | + tla::MakeShape(actualBlockShape.m(), actualBlockShape.n())); | ||
| 166 | + auto tileYT = tla::GetTile( | ||
| 167 | + tensorYT, tla::MakeCoord(blockCoord.m() * L1_TILE_M, blockCoord.n() * L1_TILE_N), | ||
| 168 | + tla::MakeShape(actualBlockShape.m(), actualBlockShape.n())); | ||
| 169 | + | ||
| 170 | + blockMmad(tileX, tileXt, tileY, tileYT, actualBlockShape, blockCoord); | ||
| 171 | + } | ||
| 172 | + | ||
| 173 | + AscendC::PipeBarrier<PIPE_ALL>(); | ||
| 174 | + } | ||
| 175 | + | ||
| 176 | + template <> | ||
| 177 | + CATLASS_DEVICE void operator()<AscendC::AIV>(Params const& /*params*/) | ||
| 178 | + {} | ||
| 179 | +}; | ||
| 180 | + | ||
| 181 | +} // namespace Catlass::Gemm::Kernel | ||
| 182 | + | ||
| 183 | + | ||
| @@ -105,6 +105,49 @@ struct CopyL0CToGmTla< | |||
| 105 | } | 105 | } |
| 106 | }; | 106 | }; |
| 107 | 107 | ||
| 108 | +template <class TensorSrc_, class ElementDst_, class LayoutDst_, class CoordDst_, bool ReluEnable_> | ||
| 109 | +struct CopyL0CToGmTla< | ||
| 110 | + Catlass::Arch::Ascend950, TensorSrc_, | ||
| 111 | + tla::Tensor<AscendC::GlobalTensor<ElementDst_>, LayoutDst_, CoordDst_, AscendC::TPosition::GM>, | ||
| 112 | + ScaleGranularity::NO_QUANT, ReluEnable_, std::enable_if_t<tla::detail::isColumnMajor<LayoutDst_>::value>> { | ||
| 113 | + using ArchTag = Catlass::Arch::Ascend950; | ||
| 114 | + using ElementDst = ElementDst_; | ||
| 115 | + using ElementSrc = typename TensorSrc_::Element; | ||
| 116 | + static constexpr auto quantPre = | ||
| 117 | + CopyL0CToDstQuantMode<ArchTag, ElementSrc, ElementDst, ScaleGranularity::NO_QUANT>::VALUE; | ||
| 118 | + static constexpr auto reluEn = ReluEnable_; | ||
| 119 | + | ||
| 120 | + template <class TensorDst, class TensorSrc> | ||
| 121 | + CATLASS_DEVICE void operator()(TensorDst const& dstTensor, TensorSrc const& srcTensor, uint8_t unitFlag = 0) | ||
| 122 | + { | ||
| 123 | + static_assert( | ||
| 124 | + tla::detail::isColumnMajor<typename TensorDst::Layout>::value && | ||
| 125 | + TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::GM, | ||
| 126 | + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be GM and ColumnMajor"); | ||
| 127 | + | ||
| 128 | + AscendC::FixpipeParamsC310<AscendC::CO2Layout::COLUMN_MAJOR> intriParams; | ||
| 129 | + | ||
| 130 | + intriParams.nSize = tla::get<1>(dstTensor.originShape()); | ||
| 131 | + intriParams.mSize = tla::get<0>(dstTensor.originShape()); | ||
| 132 | + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); | ||
| 133 | + intriParams.dstStride = tla::get<1>(dstTensor.stride()); | ||
| 134 | + intriParams.params = AscendC::Nz2DnParams( | ||
| 135 | + 1, // dnNum | ||
| 136 | + 0, // srcNzMatrixStride | ||
| 137 | + 0, // dstDnMatrixStride | ||
| 138 | + 1); // srcNzC0Stride | ||
| 139 | + intriParams.quantPre = quantPre; | ||
| 140 | + intriParams.reluEn = reluEn; | ||
| 141 | + intriParams.unitFlag = unitFlag; | ||
| 142 | + | ||
| 143 | + auto dstOffset = dstTensor.layout()(dstTensor.coord()); | ||
| 144 | + auto srcOffset = srcTensor.layout()(srcTensor.coord()); | ||
| 145 | + | ||
| 146 | + AscendC::Fixpipe<ElementDst, ElementSrc, AscendC::CFG_COLUMN_MAJOR>( | ||
| 147 | + dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); | ||
| 148 | + } | ||
| 149 | +}; | ||
| 150 | + | ||
| 108 | template <class TensorSrc_, class ElementDst_, class LayoutDst_, class CoordDst_, bool ReluEnable_> | 151 | template <class TensorSrc_, class ElementDst_, class LayoutDst_, class CoordDst_, bool ReluEnable_> |
| 109 | struct CopyL0CToGmTla< | 152 | struct CopyL0CToGmTla< |
| 110 | Catlass::Arch::Ascend950, TensorSrc_, | 153 | Catlass::Arch::Ascend950, TensorSrc_, |