已合并
feat(tile): add Ascend950 L0C-to-UB ColumnMajor specialization #1124
CheaterAbec创建于 8月18日
feat(tile): add Ascend950 L0C-to-UB ColumnMajor specialization #1124
已合并
共 1 个文件变更+83-0
| @@ -19,6 +19,7 @@ | |||
| 19 | 19 | ||
| 20 | #if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510) | 20 | #if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510) |
| 21 | constexpr AscendC::FixpipeConfig CFG_ROW_MAJOR_UB = {AscendC::CO2Layout::ROW_MAJOR, true}; | 21 | constexpr AscendC::FixpipeConfig CFG_ROW_MAJOR_UB = {AscendC::CO2Layout::ROW_MAJOR, true}; |
| 22 | +constexpr AscendC::FixpipeConfig CFG_COLUMN_MAJOR_UB = {AscendC::CO2Layout::COLUMN_MAJOR, true}; | ||
| 22 | #endif | 23 | #endif |
| 23 | 24 | ||
| 24 | namespace Catlass::Gemm::Tile { | 25 | namespace Catlass::Gemm::Tile { |
| @@ -98,6 +99,88 @@ struct CopyL0CToUBTla< | |||
| 98 | dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); | 99 | dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); |
| 99 | } | 100 | } |
| 100 | }; | 101 | }; |
| 102 | + | ||
| 103 | +template <class TensorSrc_, class ElementDst_, class LayoutDst_, class CoordDst_, bool ReluEnable_> | ||
| 104 | +struct CopyL0CToUBTla< | ||
| 105 | + Catlass::Arch::Ascend950, TensorSrc_, | ||
| 106 | + tla::Tensor<AscendC::LocalTensor<ElementDst_>, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>, | ||
| 107 | + CopyL0CToUBMode::NO_SPLIT, ScaleGranularity::NO_QUANT, ReluEnable_, | ||
| 108 | + std::enable_if_t<tla::detail::isColumnMajor<LayoutDst_>::value>> { | ||
| 109 | + using ArchTag = Catlass::Arch::Ascend950; | ||
| 110 | + using ElementDst = ElementDst_; | ||
| 111 | + using ElementSrc = typename TensorSrc_::Element; | ||
| 112 | + static constexpr auto quantPre = | ||
| 113 | + CopyL0CToDstQuantMode<ArchTag, ElementSrc, ElementDst, ScaleGranularity::NO_QUANT>::VALUE; | ||
| 114 | + static constexpr auto reluEn = ReluEnable_; | ||
| 115 | + | ||
| 116 | + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(ElementDst); | ||
| 117 | + | ||
| 118 | + template <class TensorDst, class TensorSrc> | ||
| 119 | + CATLASS_DEVICE void operator()(TensorDst const& dstTensor, TensorSrc const& srcTensor, uint8_t unitFlag = 0) | ||
| 120 | + { | ||
| 121 | + static_assert( | ||
| 122 | + tla::detail::isColumnMajor<typename TensorDst::Layout>::value && | ||
| 123 | + TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::VECCALC, | ||
| 124 | + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and ColumnMajor"); | ||
| 125 | + | ||
| 126 | + AscendC::FixpipeParamsC310<AscendC::CO2Layout::COLUMN_MAJOR> intriParams; | ||
| 127 | + | ||
| 128 | + // Fixpipe layout information | ||
| 129 | + intriParams.nSize = tla::get<1>(dstTensor.originShape()); | ||
| 130 | + intriParams.mSize = RoundUp<ELE_NUM_PER_BLK>(tla::get<0>(dstTensor.originShape())); | ||
| 131 | + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); | ||
| 132 | + intriParams.dstStride = tla::get<1>(dstTensor.stride()); | ||
| 133 | + | ||
| 134 | + intriParams.params = AscendC::Nz2DnParams(1, 0, 0, 1); // srcNzC0Stride=1 | ||
| 135 | + | ||
| 136 | + // Fixpipe auxiliary arguments | ||
| 137 | + intriParams.quantPre = quantPre; | ||
| 138 | + intriParams.reluEn = reluEn; | ||
| 139 | + intriParams.unitFlag = unitFlag; | ||
| 140 | + | ||
| 141 | + auto dstOffset = dstTensor.layout()(dstTensor.coord()); | ||
| 142 | + auto srcOffset = srcTensor.layout()(srcTensor.coord()); | ||
| 143 | + | ||
| 144 | + // Call AscendC Fixpipe | ||
| 145 | + AscendC::Fixpipe<ElementDst, ElementSrc, CFG_COLUMN_MAJOR_UB>( | ||
| 146 | + dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); | ||
| 147 | + } | ||
| 148 | + | ||
| 149 | + template <class TensorDst, class TensorSrc> | ||
| 150 | + CATLASS_DEVICE void operator()( | ||
| 151 | + TensorDst const& dstTensor, TensorSrc const& srcTensor, bool subBlockId, uint8_t unitFlag) | ||
| 152 | + { | ||
| 153 | + static_assert( | ||
| 154 | + tla::detail::isColumnMajor<typename TensorDst::Layout>::value && | ||
| 155 | + TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::VECCALC, | ||
| 156 | + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and ColumnMajor"); | ||
| 157 | + | ||
| 158 | + AscendC::FixpipeParamsC310<AscendC::CO2Layout::COLUMN_MAJOR> intriParams; | ||
| 159 | + | ||
| 160 | + // Fixpipe layout information | ||
| 161 | + intriParams.nSize = tla::get<1>(dstTensor.originShape()); | ||
| 162 | + intriParams.mSize = RoundUp<ELE_NUM_PER_BLK>(tla::get<0>(dstTensor.originShape())); | ||
| 163 | + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); | ||
| 164 | + intriParams.dstStride = tla::get<1>(dstTensor.stride()); | ||
| 165 | + | ||
| 166 | + intriParams.params = AscendC::Nz2DnParams(1, 0, 0, 1); // srcNzC0Stride=1 | ||
| 167 | + | ||
| 168 | + // Fixpipe auxiliary arguments | ||
| 169 | + intriParams.quantPre = quantPre; | ||
| 170 | + intriParams.reluEn = reluEn; | ||
| 171 | + intriParams.unitFlag = unitFlag; | ||
| 172 | + intriParams.dualDstCtl = 0; | ||
| 173 | + intriParams.subBlockId = subBlockId; | ||
| 174 | + | ||
| 175 | + auto dstOffset = dstTensor.layout()(dstTensor.coord()); | ||
| 176 | + auto srcOffset = srcTensor.layout()(srcTensor.coord()); | ||
| 177 | + | ||
| 178 | + // Call AscendC Fixpipe | ||
| 179 | + AscendC::Fixpipe<ElementDst, ElementSrc, CFG_COLUMN_MAJOR_UB>( | ||
| 180 | + dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); | ||
| 181 | + } | ||
| 182 | +}; | ||
| 183 | + | ||
| 101 | template <class TensorSrc_, class ElementDst_, class LayoutDst_, class CoordDst_, bool ReluEnable_> | 184 | template <class TensorSrc_, class ElementDst_, class LayoutDst_, class CoordDst_, bool ReluEnable_> |
| 102 | struct CopyL0CToUBTla< | 185 | struct CopyL0CToUBTla< |
| 103 | Catlass::Arch::Ascend950, TensorSrc_, | 186 | Catlass::Arch::Ascend950, TensorSrc_, |