已合并
fix CV epilogue stage row sizing #10799
yangzw653创建于 1 天前
fix CV epilogue stage row sizing #10799
已合并
共 1 个文件变更+8-6
| @@ -63,8 +63,9 @@ public: | |||
| 63 | 63 | ||
| 64 | // GM ADDR | 64 | // GM ADDR |
| 65 | AscendC::LocalTensor<DataTypeIn> cLocal_{AscendC::TPosition::VECIN, 0, AscendC::TOTAL_UB_SIZE}; | 65 | AscendC::LocalTensor<DataTypeIn> cLocal_{AscendC::TPosition::VECIN, 0, AscendC::TOTAL_UB_SIZE}; |
| 66 | - // vector核一次最多计算多少个元素 | 66 | + // AF会融合where、compare等算子, 其数据位宽比matmul输出类型小, |
| 67 | - int64_t stageSize_ = 0; | 67 | + // 若stage仍按元素个数切分会无法统一vector核UB数据计算, 因此统一到矩阵行: M轴一次最多处理的行数 |
| 68 | + int64_t cubeMStageRows_ = 0; | ||
| 68 | // attribute | 69 | // attribute |
| 69 | FusionOp fusionOp_; | 70 | FusionOp fusionOp_; |
| 70 | ProblemShape problemShape_; | 71 | ProblemShape problemShape_; |
| @@ -73,8 +74,8 @@ public: | |||
| 73 | { | 74 | { |
| 74 | int64_t l1NAlign = AlignBlock<DataTypeOut>(l1N); | 75 | int64_t l1NAlign = AlignBlock<DataTypeOut>(l1N); |
| 75 | int64_t ubOffset = l1M * l1NAlign; | 76 | int64_t ubOffset = l1M * l1NAlign; |
| 76 | - // 基于剩余UB可用大小确定stageSize_ | 77 | + // 基于剩余UB可用大小确定cubeMStageRows_(元素个数) |
| 77 | - fusionOp_.Init(params.fusionParams, cLocal_, l1M, l1NAlign, ubOffset, stageSize_); | 78 | + fusionOp_.Init(params.fusionParams, cLocal_, l1M, l1NAlign, ubOffset, cubeMStageRows_); |
| 78 | problemShape_ = problemShape; | 79 | problemShape_ = problemShape; |
| 79 | } | 80 | } |
| 80 | 81 | ||
| @@ -90,8 +91,9 @@ public: | |||
| 90 | int64_t blockShapeNAlign = AlignBlock<DataTypeOut>(blockShapeN); // 对齐16 | 91 | int64_t blockShapeNAlign = AlignBlock<DataTypeOut>(blockShapeN); // 对齐16 |
| 91 | int64_t inputSize = blockShapeM * blockShapeNAlign; | 92 | int64_t inputSize = blockShapeM * blockShapeNAlign; |
| 92 | 93 | ||
| 93 | - // 一次计算最多取Min(baseM/2 * baseN, stageSize_) | 94 | + // 每个stage统一按矩阵行处理: stageM取Min(cubeMStageRows_, blockShapeM), stageSize = stageM * blockShapeNAlign |
| 94 | - int64_t stageSize = AscendC::Std::min(stageSize_, inputSize) / blockShapeNAlign * blockShapeNAlign; | 95 | + const int64_t stageM = AscendC::Std::min(cubeMStageRows_, blockShapeM); |
| 96 | + int64_t stageSize = stageM * blockShapeNAlign; | ||
| 95 | // m轴为1场景第二个vec直接返回 | 97 | // m轴为1场景第二个vec直接返回 |
| 96 | if (stageSize <= 0) { | 98 | if (stageSize <= 0) { |
| 97 | AscendC::CrossCoreSetFlag<AIC_SYNC_AIV_MODE_4, PIPE_V>(flagId); | 99 | AscendC::CrossCoreSetFlag<AIC_SYNC_AIV_MODE_4, PIPE_V>(flagId); |