// // FmhaV2Execution.cpp // MNN // // Created by MNN on 2024/06/03. // Copyright © 2018, Alibaba Group Holding Limited // #ifdef MNN_SUPPORT_TRANSFORMER_FUSE #include #include #include "backend/opencl/execution/buffer/SelfAttentionBufExecution.hpp" namespace MNN { namespace OpenCL { SelfAttentionBufImpl::SelfAttentionBufImpl(const MNN::Op *op, Backend *backend){ auto fmha_v2_param = op->main_as_FmhaV2Param(); mNumHead = fmha_v2_param->heads(); mOpenCLBackend = static_cast(backend); auto kernel = mOpenCLBackend->getOpenCLRuntime()->buildKernel("self_attention_buf", "softmax_inside", {"-DSOFTMAX_LOCAL_SIZE=512"}, mOpenCLBackend->getPrecision()); OPENCL_CHECK_KERNEL_CTOR(kernel); mMaxWorkGroupSize = static_cast(mOpenCLBackend->getOpenCLRuntime()->getMaxWorkGroupSize(kernel)); } int SelfAttentionBufImpl::getLocalSize(int size, int maxGroupSize){ int local_size = 1; while(local_size * 2 <= maxGroupSize && local_size * 2 <= size){ local_size *= 2; } return local_size; } // [B, seqlen, HeadNum*3*HeadDim] -> [B, seqlen, HeadNum*HeadDim] ErrorCode SelfAttentionBufImpl::onResize(Backend *backend, const std::vector &inputs, const std::vector &outputs) { mOpenCLBackend = static_cast(backend); mOpenCLBackend->startRecord(mRecording); auto input = inputs[0];// [Batch, seqLen, mNumHead * 3 * mHeadDim] auto runtime = mOpenCLBackend->getOpenCLRuntime(); auto shape = input->shape(); int tile_mn = 32; int tile_k = 4; // for gemm alignment int batch = shape[0]; int seq_len = shape[1]; mHeadDim = shape[2] / mNumHead / 3; mScale = 1.0 / sqrt(mHeadDim); if(mOpenCLBackend->getPrecision() != BackendConfig::Precision_High){ mByte = 2; } // split several pieces for memory save if(seq_len > 1024) { mQseqSplitNum = (seq_len >= 4096 && seq_len % 64 == 0) ? 8 : ((seq_len < 2048) ? 2 : 4); } // splitPiecesSize need aligned to 32, make sure XgemmBatched globalsize be divisible by localsize int splitPiecesSize = ROUND_UP(seq_len, tile_mn) / mQseqSplitNum; while((splitPiecesSize % 32) != 0){ tile_mn *= 2; splitPiecesSize = ROUND_UP(seq_len, tile_mn) / mQseqSplitNum; } int buffer_size = batch * mNumHead * ROUND_UP(mHeadDim, tile_k) * ROUND_UP(seq_len, tile_mn); int buffer_qk_size = batch * mNumHead * ROUND_UP(seq_len, tile_mn) * ROUND_UP(seq_len, tile_mn) / mQseqSplitNum; int buffer_v_size = batch * mNumHead * ROUND_UP(mHeadDim, tile_mn) * ROUND_UP(seq_len, tile_mn); mTempQ.reset(Tensor::createDevice(std::vector{buffer_size / mQseqSplitNum})); mTempK.reset(Tensor::createDevice(std::vector{buffer_size})); mTempV.reset(Tensor::createDevice(std::vector{buffer_v_size})); mTempQK.reset(Tensor::createDevice(std::vector{buffer_qk_size})); mTempTrans.reset(Tensor::createDevice(std::vector{buffer_qk_size})); mTempSoftMax.reset(Tensor::createDevice(std::vector{buffer_qk_size})); mTempQKV.reset(Tensor::createDevice(std::vector{buffer_v_size / mQseqSplitNum})); // printf("buffer size x2:%f MB, buffer qk size x3:%f MB, buffer v size x2 :%f MB\n", buffer_size * 2.0 / 1024.0 / 1024.0, buffer_qk_size * 2.0 / 1024.0 / 1024.0, buffer_v_size * 2.0 / 1024.0 / 1024.0); mOpenCLBackend->onAcquireBuffer(mTempQ.get(), Backend::DYNAMIC); mOpenCLBackend->onAcquireBuffer(mTempK.get(), Backend::DYNAMIC); mOpenCLBackend->onAcquireBuffer(mTempV.get(), Backend::DYNAMIC); mOpenCLBackend->onAcquireBuffer(mTempQK.get(), Backend::DYNAMIC); mOpenCLBackend->onReleaseBuffer(mTempQ.get(), Backend::DYNAMIC); mOpenCLBackend->onReleaseBuffer(mTempK.get(), Backend::DYNAMIC); mOpenCLBackend->onAcquireBuffer(mTempSoftMax.get(), Backend::DYNAMIC); mOpenCLBackend->onReleaseBuffer(mTempQK.get(), Backend::DYNAMIC); mOpenCLBackend->onAcquireBuffer(mTempTrans.get(), Backend::DYNAMIC); mOpenCLBackend->onReleaseBuffer(mTempSoftMax.get(), Backend::DYNAMIC); mOpenCLBackend->onAcquireBuffer(mTempQKV.get(), Backend::DYNAMIC); mOpenCLBackend->onReleaseBuffer(mTempV.get(), Backend::DYNAMIC); mOpenCLBackend->onReleaseBuffer(mTempTrans.get(), Backend::DYNAMIC); mOpenCLBackend->onReleaseBuffer(mTempQKV.get(), Backend::DYNAMIC); mKernel_split.resize(mQseqSplitNum); mKernel_qk.resize(mQseqSplitNum); mKernel_softmax.resize(mQseqSplitNum); mKernel_qkv.resize(mQseqSplitNum); mKernel_clip.resize(mQseqSplitNum); mKernel_trans.resize(mQseqSplitNum); mGlobalWorkSizeSplit.resize(mQseqSplitNum); mLocalWorkSizeSplit.resize(mQseqSplitNum); mGlobalWorkSizeClip.resize(mQseqSplitNum); mLocalWorkSizeClip.resize(mQseqSplitNum); mGlobalWorkSizeQk.resize(mQseqSplitNum); mLocalWorkSizeQk.resize(mQseqSplitNum); mGlobalWorkSizeSoftMax.resize(mQseqSplitNum); mLocalWorkSizeSoftMax.resize(mQseqSplitNum); mGlobalWorkSizeQkv.resize(mQseqSplitNum); mLocalWorkSizeQkv.resize(mQseqSplitNum); mGlobalWorkSizeTrans.resize(mQseqSplitNum); mLocalWorkSizeTrans.resize(mQseqSplitNum); for(int seq_idx = 0; seq_idx < mQseqSplitNum; seq_idx++) { // Split input to q k v { // [Batch, seqLen, mNumHead * 3 * mHeadDim] -> // Q : [Batch * mNumHead, ROUND_UP(mHeadDim, tile_k), ROUND_UP(seqLen, tile_mn)] // K : [Batch * mNumHead, ROUND_UP(mHeadDim, tile_k), ROUND_UP(seqLen, tile_mn)] // V : [Batch * mNumHead, ROUND_UP(seqLen, tile_mn), ROUND_UP(mHeadDim, tile_mn)] std::set buildOption; if((mHeadDim % 4) != 0){ buildOption.emplace("-DHEADDIM_LEAVE"); } if((seq_len % 4) != 0){ buildOption.emplace("-DSEQLEN_LEAVE"); } int seq_len_pack_mn = ROUND_UP(seq_len, tile_mn); int head_dim_pack_mn = ROUND_UP(mHeadDim, tile_mn); int head_dim_pack_k = ROUND_UP(mHeadDim, tile_k); int seq_len_piece = seq_len_pack_mn/mQseqSplitNum; mKernel_split[seq_idx] = runtime->buildKernel("self_attention_buf", "split_transpose_qkv", buildOption, mOpenCLBackend->getPrecision(), inputs[0], outputs[0]); auto maxWorkGroupSize = static_cast(runtime->getMaxWorkGroupSize(mKernel_split[seq_idx])); mGlobalWorkSizeSplit[seq_idx] = {static_cast(UP_DIV(seq_len_pack_mn, 4)), static_cast(UP_DIV(head_dim_pack_mn, 4)), static_cast(batch*mNumHead)}; if(seq_idx > 0) { mGlobalWorkSizeSplit[seq_idx][0] = static_cast(UP_DIV(seq_len_piece, 4)); } uint32_t index = 0; cl_int ret = CL_SUCCESS; ret |= mKernel_split[seq_idx]->get().setArg(index++, mGlobalWorkSizeSplit[seq_idx][0]); ret |= mKernel_split[seq_idx]->get().setArg(index++, mGlobalWorkSizeSplit[seq_idx][1]); ret |= mKernel_split[seq_idx]->get().setArg(index++, mGlobalWorkSizeSplit[seq_idx][2]); ret |= mKernel_split[seq_idx]->get().setArg(index++, openCLBuffer(input)); ret |= mKernel_split[seq_idx]->get().setArg(index++, openCLBuffer(mTempQ.get())); ret |= mKernel_split[seq_idx]->get().setArg(index++, openCLBuffer(mTempK.get())); ret |= mKernel_split[seq_idx]->get().setArg(index++, openCLBuffer(mTempV.get())); ret |= mKernel_split[seq_idx]->get().setArg(index++, seq_len_pack_mn); ret |= mKernel_split[seq_idx]->get().setArg(index++, seq_len_piece); ret |= mKernel_split[seq_idx]->get().setArg(index++, head_dim_pack_mn); ret |= mKernel_split[seq_idx]->get().setArg(index++, head_dim_pack_k); ret |= mKernel_split[seq_idx]->get().setArg(index++, seq_len); ret |= mKernel_split[seq_idx]->get().setArg(index++, mNumHead); ret |= mKernel_split[seq_idx]->get().setArg(index++, mHeadDim); ret |= mKernel_split[seq_idx]->get().setArg(index++, batch); ret |= mKernel_split[seq_idx]->get().setArg(index++, seq_idx); MNN_CHECK_CL_SUCCESS(ret, "setArg split_transpose_qkv"); mLocalWorkSizeSplit[seq_idx] = localWS3DDefault(mGlobalWorkSizeSplit[seq_idx], maxWorkGroupSize, runtime, "split_transpose_qkv", mKernel_split[seq_idx], mOpenCLBackend->getCLTuneLevel(), "self_attention_buf").first; mGlobalWorkSizeSplit[seq_idx][0] = ROUND_UP(mGlobalWorkSizeSplit[seq_idx][0], std::max((uint32_t)1, mLocalWorkSizeSplit[seq_idx][0])); mGlobalWorkSizeSplit[seq_idx][1] = ROUND_UP(mGlobalWorkSizeSplit[seq_idx][1], std::max((uint32_t)1, mLocalWorkSizeSplit[seq_idx][1])); mGlobalWorkSizeSplit[seq_idx][2] = ROUND_UP(mGlobalWorkSizeSplit[seq_idx][2], std::max((uint32_t)1, mLocalWorkSizeSplit[seq_idx][2])); mOpenCLBackend->recordKernel3d(mKernel_split[seq_idx], mGlobalWorkSizeSplit[seq_idx], mLocalWorkSizeSplit[seq_idx]); } // query * key -> div { // Q : [Batch * mNumHead, ROUND_UP(mHeadDim, tile_k), ROUND_UP(seqLen, tile_mn)] -> [B, K, M] // K : [Batch * mNumHead, ROUND_UP(mHeadDim, tile_k), ROUND_UP(seqLen, tile_mn)] -> [B, K, N] // QV: [Batch * mNumHead, ROUND_UP(seqLen, tile_mn), ROUND_UP(seqLen, tile_mn)] -> [B, N, M] int loop = batch * mNumHead; int e_pack = ROUND_UP(seq_len, tile_mn) / mQseqSplitNum; int l_pack = ROUND_UP(mHeadDim, tile_k); int h_pack = ROUND_UP(seq_len, tile_mn); std::set buildOptions; uint32_t layout = 4; auto param = getGemmParams({(uint32_t)e_pack, (uint32_t)h_pack, (uint32_t)l_pack, layout, (uint32_t)loop, (uint32_t)0}, {openCLBuffer(mTempQ.get()), openCLBuffer(mTempK.get()), openCLBuffer(mTempQK.get())}, mOpenCLBackend->getOpenCLRuntime(), mOpenCLBackend->getPrecision(), mOpenCLBackend->getCLTuneLevel()); int KWG=param[0], KWI=param[1], MDIMA=param[2], MDIMC=param[3], MWG=param[4], NDIMB=param[5], NDIMC=param[6], NWG=param[7], SA=param[8], SB=param[9], STRM=param[10], STRN=param[11], VWM=param[12], VWN=param[13]; buildOptions.emplace("-DKWG=" + std::to_string(KWG)); buildOptions.emplace("-DKWI=" + std::to_string(KWI)); buildOptions.emplace("-DMDIMA=" + std::to_string(MDIMA)); buildOptions.emplace("-DMDIMC=" + std::to_string(MDIMC)); buildOptions.emplace("-DMWG=" + std::to_string(MWG)); buildOptions.emplace("-DNDIMB=" + std::to_string(NDIMB)); buildOptions.emplace("-DNDIMC=" + std::to_string(NDIMC)); buildOptions.emplace("-DNWG=" + std::to_string(NWG)); buildOptions.emplace("-DSA=" + std::to_string(SA)); buildOptions.emplace("-DSB=" + std::to_string(SB)); buildOptions.emplace("-DSTRM=" + std::to_string(STRM)); buildOptions.emplace("-DSTRN=" + std::to_string(STRN)); buildOptions.emplace("-DVWM=" + std::to_string(VWM)); buildOptions.emplace("-DVWN=" + std::to_string(VWN)); if(layout >= 4) { buildOptions.emplace("-DOUTPUTMN"); } int tileM = MWG; int tileN = NWG; int localM = MDIMC; int localN = NDIMC; if(mOpenCLBackend->getOpenCLRuntime()->getGpuType() == GpuType::ADRENO) { buildOptions.emplace("-DUSE_CL_MAD=1"); buildOptions.emplace("-DRELAX_WORKGROUP_SIZE=1"); } buildOptions.emplace("-DONLY_HAVE_ALPHA"); mKernel_qk[seq_idx] = mOpenCLBackend->getOpenCLRuntime()->buildKernel("matmul_params_buf", "XgemmBatched", buildOptions, mOpenCLBackend->getPrecision()); int out_per_thread_m = tileM / localM; int out_per_thread_n = tileN / localN; mGlobalWorkSizeQk[seq_idx] = {static_cast(e_pack/out_per_thread_m), static_cast(h_pack/out_per_thread_n), static_cast(loop)}; mLocalWorkSizeQk[seq_idx] = {static_cast(localM), static_cast(localN), 1}; float alpha = mScale; float beta = 0.0f; int batch_offset_a = e_pack * l_pack; int batch_offset_b = h_pack * l_pack; int batch_offset_c = e_pack * h_pack; int batch_offset[4] = {batch_offset_a, batch_offset_b, batch_offset_c, 0}; int base_ptr_offset[4] = {0, 0, 0, 0}; int stride[4] = {e_pack, h_pack, h_pack, h_pack}; int group[4] = {1, 1, 1, loop}; int idx = 0; cl_int ret = CL_SUCCESS; ret |= mKernel_qk[seq_idx]->get().setArg(idx++, static_cast(e_pack)); ret |= mKernel_qk[seq_idx]->get().setArg(idx++, static_cast(h_pack)); ret |= mKernel_qk[seq_idx]->get().setArg(idx++, static_cast(l_pack)); ret |= mKernel_qk[seq_idx]->get().setArg(idx++, alpha); ret |= mKernel_qk[seq_idx]->get().setArg(idx++, beta); ret |= mKernel_qk[seq_idx]->get().setArg(idx++, openCLBuffer(mTempQ.get())); ret |= mKernel_qk[seq_idx]->get().setArg(idx++, openCLBuffer(mTempK.get())); ret |= mKernel_qk[seq_idx]->get().setArg(idx++, openCLBuffer(mTempQK.get())); ret |= mKernel_qk[seq_idx]->get().setArg(idx++, batch_offset); ret |= mKernel_qk[seq_idx]->get().setArg(idx++, base_ptr_offset); ret |= mKernel_qk[seq_idx]->get().setArg(idx++, stride); ret |= mKernel_qk[seq_idx]->get().setArg(idx++, group); MNN_CHECK_CL_SUCCESS(ret, "setArg Self-Attention batchmatmul qk Kernel"); mOpenCLBackend->recordKernel3d(mKernel_qk[seq_idx], mGlobalWorkSizeQk[seq_idx], mLocalWorkSizeQk[seq_idx]); } // softmax { // QV: [Batch * mNumHead, ROUND_UP(seqLen, tile_mn), ROUND_UP(seqLen, tile_mn)] // Sotmax: [Batch * mNumHead, ROUND_UP(seqLen, tile_mn), ROUND_UP(seqLen, tile_mn)] // axis : 1 (middle dim) mSoftmaxShape[0] = batch*mNumHead; mSoftmaxShape[1] = ROUND_UP(seq_len, tile_mn)/mQseqSplitNum; mSoftmaxShape[2] = ROUND_UP(seq_len, tile_mn); auto MaxLocalSize = std::min(std::min(runtime->getMaxWorkItemSizes()[0], mMaxWorkGroupSize), static_cast(256)); int localSize = getLocalSize(mSoftmaxShape[1], MaxLocalSize); if(localSize < 4){ localSize = 1; } std::set buildOption; buildOption.emplace("-DSOFTMAX_LOCAL_SIZE=" + std::to_string(localSize)); // buildOption.emplace("-DOUTPUT_TRANSPOSE"); mKernel_softmax[seq_idx] = runtime->buildKernel("self_attention_buf", "softmax_inside", buildOption, mOpenCLBackend->getPrecision(), inputs[0], outputs[0]); mGlobalWorkSizeSoftMax[seq_idx] = {static_cast(localSize), static_cast(mSoftmaxShape[1]), static_cast(mSoftmaxShape[0])}; uint32_t index = 0; cl_int ret = CL_SUCCESS; ret |= mKernel_softmax[seq_idx]->get().setArg(index++, mGlobalWorkSizeSoftMax[seq_idx][0]); ret |= mKernel_softmax[seq_idx]->get().setArg(index++, mGlobalWorkSizeSoftMax[seq_idx][1]); ret |= mKernel_softmax[seq_idx]->get().setArg(index++, mGlobalWorkSizeSoftMax[seq_idx][2]); ret |= mKernel_softmax[seq_idx]->get().setArg(index++, openCLBuffer(mTempQK.get())); ret |= mKernel_softmax[seq_idx]->get().setArg(index++, openCLBuffer(mTempSoftMax.get())); ret |= mKernel_softmax[seq_idx]->get().setArg(index++, seq_len); ret |= mKernel_softmax[seq_idx]->get().setArg(index++, mSoftmaxShape); MNN_CHECK_CL_SUCCESS(ret, "setArg Self-Attention softmax"); mLocalWorkSizeSoftMax[seq_idx] = {static_cast(localSize), 1, 1}; mOpenCLBackend->recordKernel3d(mKernel_softmax[seq_idx], mGlobalWorkSizeSoftMax[seq_idx], mLocalWorkSizeSoftMax[seq_idx]); } { int loop = batch * mNumHead; int transDimW = ROUND_UP(seq_len, tile_mn) / mQseqSplitNum; int transDimH = ROUND_UP(seq_len, tile_mn); std::set buildOptions; mKernel_trans[seq_idx] = runtime->buildKernel("self_attention_buf", "trans_3d_buf", buildOptions, mOpenCLBackend->getPrecision(), inputs[0], outputs[0]); uint32_t maxWorkGroupSize = static_cast(mOpenCLBackend->getOpenCLRuntime()->getMaxWorkGroupSize(mKernel_trans[seq_idx])); mGlobalWorkSizeTrans[seq_idx] = {(uint32_t)transDimW/8, (uint32_t)transDimH/8, (uint32_t)(loop)}; uint32_t index = 0; cl_int ret = CL_SUCCESS; ret |= mKernel_trans[seq_idx]->get().setArg(index++, mGlobalWorkSizeTrans[seq_idx][0]); ret |= mKernel_trans[seq_idx]->get().setArg(index++, mGlobalWorkSizeTrans[seq_idx][1]); ret |= mKernel_trans[seq_idx]->get().setArg(index++, mGlobalWorkSizeTrans[seq_idx][2]); ret |= mKernel_trans[seq_idx]->get().setArg(index++, openCLBuffer(mTempSoftMax.get())); ret |= mKernel_trans[seq_idx]->get().setArg(index++, openCLBuffer(mTempTrans.get())); ret |= mKernel_trans[seq_idx]->get().setArg(index++, loop); ret |= mKernel_trans[seq_idx]->get().setArg(index++, transDimW); ret |= mKernel_trans[seq_idx]->get().setArg(index++, transDimH); MNN_CHECK_CL_SUCCESS(ret, "setArg Self-Attention transpose"); mLocalWorkSizeTrans[seq_idx] = localWS3DDefault(mGlobalWorkSizeTrans[seq_idx], maxWorkGroupSize, mOpenCLBackend->getOpenCLRuntime(), "trans_3d_buf", mKernel_trans[seq_idx], mOpenCLBackend->getCLTuneLevel(), "self_attention_buf").first; mGlobalWorkSizeTrans[seq_idx][0] = ROUND_UP(mGlobalWorkSizeTrans[seq_idx][0], std::max((uint32_t)1, mLocalWorkSizeTrans[seq_idx][0])); mGlobalWorkSizeTrans[seq_idx][1] = ROUND_UP(mGlobalWorkSizeTrans[seq_idx][1], std::max((uint32_t)1, mLocalWorkSizeTrans[seq_idx][1])); mGlobalWorkSizeTrans[seq_idx][2] = ROUND_UP(mGlobalWorkSizeTrans[seq_idx][2], std::max((uint32_t)1, mLocalWorkSizeTrans[seq_idx][2])); mOpenCLBackend->recordKernel3d(mKernel_trans[seq_idx], mGlobalWorkSizeTrans[seq_idx], mLocalWorkSizeTrans[seq_idx]); } // qk * value { // Sotmax: [Batch * mNumHead, ROUND_UP(seqLen, tile), ROUND_UP(seqLen, tile)] -> [B, K, M] // V : [Batch * mNumHead, ROUND_UP(seqLen, tile), ROUND_UP(mHeadDim, tile)] -> [B, K, N] // QKV : [Batch * mNumHead, ROUND_UP(mHeadDim, tile), ROUND_UP(seqLen, tile)] -> [B, N, M] int loop = batch * mNumHead; int e_pack = ROUND_UP(seq_len, tile_mn) / mQseqSplitNum; int l_pack = ROUND_UP(seq_len, tile_mn); int h_pack = ROUND_UP(mHeadDim, tile_mn); std::set buildOptions; /* 0 -> A:[K, M] B:[K, N] C:[N, M] 1 -> A:[K, M] B:[N, K] C:[N, M] 2 -> A:[M, K] B:[K, N] C:[N, M] 3 -> A:[M, K] B:[N, K] C:[N, M] 4 -> A:[K, M] B:[K, N] C:[M, N] 5 -> A:[K, M] B:[N, K] C:[M, N] 6 -> A:[M, K] B:[K, N] C:[M, N] 7 -> A:[M, K] B:[N, K] C:[M, N] */ uint32_t layout = 0; auto param = getGemmParams({(uint32_t)e_pack, (uint32_t)h_pack, (uint32_t)l_pack, layout, (uint32_t)loop, (uint32_t)0}, {openCLBuffer(mTempTrans.get()), openCLBuffer(mTempV.get()), openCLBuffer(mTempQKV.get())}, mOpenCLBackend->getOpenCLRuntime(), mOpenCLBackend->getPrecision(), mOpenCLBackend->getCLTuneLevel()); int KWG=param[0], KWI=param[1], MDIMA=param[2], MDIMC=param[3], MWG=param[4], NDIMB=param[5], NDIMC=param[6], NWG=param[7], SA=param[8], SB=param[9], STRM=param[10], STRN=param[11], VWM=param[12], VWN=param[13]; buildOptions.emplace("-DKWG=" + std::to_string(KWG)); buildOptions.emplace("-DKWI=" + std::to_string(KWI)); buildOptions.emplace("-DMDIMA=" + std::to_string(MDIMA)); buildOptions.emplace("-DMDIMC=" + std::to_string(MDIMC)); buildOptions.emplace("-DMWG=" + std::to_string(MWG)); buildOptions.emplace("-DNDIMB=" + std::to_string(NDIMB)); buildOptions.emplace("-DNDIMC=" + std::to_string(NDIMC)); buildOptions.emplace("-DNWG=" + std::to_string(NWG)); buildOptions.emplace("-DSA=" + std::to_string(SA)); buildOptions.emplace("-DSB=" + std::to_string(SB)); buildOptions.emplace("-DSTRM=" + std::to_string(STRM)); buildOptions.emplace("-DSTRN=" + std::to_string(STRN)); buildOptions.emplace("-DVWM=" + std::to_string(VWM)); buildOptions.emplace("-DVWN=" + std::to_string(VWN)); if(layout >= 4) { buildOptions.emplace("-DOUTPUTMN"); } int tileM = MWG; int tileN = NWG; int localM = MDIMC; int localN = NDIMC; if(mOpenCLBackend->getOpenCLRuntime()->getGpuType() == GpuType::ADRENO) { buildOptions.emplace("-DUSE_CL_MAD=1"); buildOptions.emplace("-DRELAX_WORKGROUP_SIZE=1"); } mKernel_qkv[seq_idx] = mOpenCLBackend->getOpenCLRuntime()->buildKernel("matmul_params_buf", "XgemmBatched", buildOptions, mOpenCLBackend->getPrecision()); int out_per_thread_m = tileM / localM; int out_per_thread_n = tileN / localN; mGlobalWorkSizeQkv[seq_idx] = {static_cast(e_pack/out_per_thread_m), static_cast(h_pack/out_per_thread_n), static_cast(loop)}; mLocalWorkSizeQkv[seq_idx] = {static_cast(localM), static_cast(localN), 1}; float alpha = 1.0f; float beta = 0.0f; int batch_offset_a = e_pack * l_pack; int batch_offset_b = h_pack * l_pack; int batch_offset_c = e_pack * h_pack; int batch_offset[4] = {batch_offset_a, batch_offset_b, batch_offset_c, 0}; int base_ptr_offset[4] = {0, 0, 0, 0}; int stride[4] = {e_pack, h_pack, e_pack, h_pack}; int group[4] = {1, 1, 1, loop}; int idx = 0; cl_int ret = CL_SUCCESS; ret |= mKernel_qkv[seq_idx]->get().setArg(idx++, static_cast(e_pack)); ret |= mKernel_qkv[seq_idx]->get().setArg(idx++, static_cast(h_pack)); ret |= mKernel_qkv[seq_idx]->get().setArg(idx++, static_cast(l_pack)); ret |= mKernel_qkv[seq_idx]->get().setArg(idx++, alpha); ret |= mKernel_qkv[seq_idx]->get().setArg(idx++, beta); ret |= mKernel_qkv[seq_idx]->get().setArg(idx++, openCLBuffer(mTempTrans.get())); ret |= mKernel_qkv[seq_idx]->get().setArg(idx++, openCLBuffer(mTempV.get())); ret |= mKernel_qkv[seq_idx]->get().setArg(idx++, openCLBuffer(mTempQKV.get())); ret |= mKernel_qkv[seq_idx]->get().setArg(idx++, batch_offset); ret |= mKernel_qkv[seq_idx]->get().setArg(idx++, base_ptr_offset); ret |= mKernel_qkv[seq_idx]->get().setArg(idx++, stride); ret |= mKernel_qkv[seq_idx]->get().setArg(idx++, group); MNN_CHECK_CL_SUCCESS(ret, "setArg Self-Attention batchmatmul qkv Kernel"); mOpenCLBackend->recordKernel3d(mKernel_qkv[seq_idx], mGlobalWorkSizeQkv[seq_idx], mLocalWorkSizeQkv[seq_idx]); } // transpose to output { // QKV : [Batch * mNumHead, ROUND_UP(mHeadDim, tile_mn), ROUND_UP(seqLen, tile_mn)] -> [B, N, M] // output: [Batch, seqLen, mNumHead * mHeadDim] std::set buildOption; mKernel_clip[seq_idx] = runtime->buildKernel("self_attention_buf", "clip_transpose_qkv", buildOption, mOpenCLBackend->getPrecision(), inputs[0], outputs[0]); auto maxWorkGroupSize = static_cast(runtime->getMaxWorkGroupSize(mKernel_clip[seq_idx])); int seq_len_piece = ROUND_UP(seq_len, tile_mn) / mQseqSplitNum; mGlobalWorkSizeClip[seq_idx] = {static_cast(UP_DIV(seq_len_piece, 4)), static_cast(UP_DIV(mHeadDim, 4)), static_cast(batch*mNumHead)}; uint32_t index = 0; cl_int ret = CL_SUCCESS; ret |= mKernel_clip[seq_idx]->get().setArg(index++, mGlobalWorkSizeClip[seq_idx][0]); ret |= mKernel_clip[seq_idx]->get().setArg(index++, mGlobalWorkSizeClip[seq_idx][1]); ret |= mKernel_clip[seq_idx]->get().setArg(index++, mGlobalWorkSizeClip[seq_idx][2]); ret |= mKernel_clip[seq_idx]->get().setArg(index++, openCLBuffer(mTempQKV.get())); ret |= mKernel_clip[seq_idx]->get().setArg(index++, openCLBuffer(outputs[0])); ret |= mKernel_clip[seq_idx]->get().setArg(index++, tile_mn); ret |= mKernel_clip[seq_idx]->get().setArg(index++, seq_len); ret |= mKernel_clip[seq_idx]->get().setArg(index++, seq_len_piece); ret |= mKernel_clip[seq_idx]->get().setArg(index++, mNumHead); ret |= mKernel_clip[seq_idx]->get().setArg(index++, mHeadDim); ret |= mKernel_clip[seq_idx]->get().setArg(index++, batch); ret |= mKernel_clip[seq_idx]->get().setArg(index++, seq_idx); mLocalWorkSizeClip[seq_idx] = localWS3DDefault(mGlobalWorkSizeClip[seq_idx], maxWorkGroupSize, runtime, "clip_transpose_qkv", mKernel_clip[seq_idx], mOpenCLBackend->getCLTuneLevel(), "self_attention_buf").first; mGlobalWorkSizeClip[seq_idx][0] = ROUND_UP(mGlobalWorkSizeClip[seq_idx][0], std::max((uint32_t)1, mLocalWorkSizeClip[seq_idx][0])); mGlobalWorkSizeClip[seq_idx][1] = ROUND_UP(mGlobalWorkSizeClip[seq_idx][1], std::max((uint32_t)1, mLocalWorkSizeClip[seq_idx][1])); mGlobalWorkSizeClip[seq_idx][2] = ROUND_UP(mGlobalWorkSizeClip[seq_idx][2], std::max((uint32_t)1, mLocalWorkSizeClip[seq_idx][2])); MNN_CHECK_CL_SUCCESS(ret, "setArg clip_transpose_qkv"); mOpenCLBackend->recordKernel3d(mKernel_clip[seq_idx], mGlobalWorkSizeClip[seq_idx], mLocalWorkSizeClip[seq_idx]); } } mOpenCLBackend->endRecord(mRecording); return NO_ERROR; } ErrorCode SelfAttentionBufImpl::onExecute(Backend *backend, const std::vector &inputs, const std::vector &outputs) { #ifdef LOG_VERBOSE MNN_PRINT("start SelfAttentionBufExecution onExecute !\n"); #endif mOpenCLBackend = static_cast(backend); #ifdef ENABLE_OPENCL_TIME_PROFILER int batch = inputs[0]->shape()[0]; int seqLen = inputs[0]->shape()[1]; int headDim = inputs[0]->shape()[2]/3/mNumHead; std::string name; name += "-b" + std::to_string(batch); name += "-s" + std::to_string(seqLen); name += "-h" + std::to_string(mNumHead); name += "-d" + std::to_string(headDim); for(int seq_idx = 0; seq_idx < mQseqSplitNum; seq_idx++) { { cl::Event event; run3DKernelDefault(mKernel_split[seq_idx], mGlobalWorkSizeSplit[seq_idx], mLocalWorkSizeSplit[seq_idx], mOpenCLBackend->getOpenCLRuntime(), &event); mOpenCLBackend->getOpenCLRuntime()->pushEvent({"While-gemm-split" + name, event}); } { cl::Event event; run3DKernelDefault(mKernel_qk[seq_idx], mGlobalWorkSizeQk[seq_idx], mLocalWorkSizeQk[seq_idx], mOpenCLBackend->getOpenCLRuntime(), &event); mOpenCLBackend->getOpenCLRuntime()->pushEvent({"While-gemm-batchgemm" + name, event}); } { cl::Event event; run3DKernelDefault(mKernel_softmax[seq_idx], mGlobalWorkSizeSoftMax[seq_idx], mLocalWorkSizeSoftMax[seq_idx], mOpenCLBackend->getOpenCLRuntime(), &event); mOpenCLBackend->getOpenCLRuntime()->pushEvent({"While-gemm-softmax" + name, event}); } { cl::Event event; run3DKernelDefault(mKernel_trans[seq_idx], mGlobalWorkSizeTrans[seq_idx], mLocalWorkSizeTrans[seq_idx], mOpenCLBackend->getOpenCLRuntime(), &event); mOpenCLBackend->getOpenCLRuntime()->pushEvent({"While-gemm-trans-1" + name, event}); } { cl::Event event; run3DKernelDefault(mKernel_qkv[seq_idx], mGlobalWorkSizeQkv[seq_idx], mLocalWorkSizeQkv[seq_idx], mOpenCLBackend->getOpenCLRuntime(), &event); mOpenCLBackend->getOpenCLRuntime()->pushEvent({"While-gemm-batchgemm" + name, event}); } { cl::Event event; run3DKernelDefault(mKernel_clip[seq_idx], mGlobalWorkSizeClip[seq_idx], mLocalWorkSizeClip[seq_idx], mOpenCLBackend->getOpenCLRuntime(), &event); mOpenCLBackend->getOpenCLRuntime()->pushEvent({"While-gemm-clip" + name, event}); } } #else for(int seq_idx = 0; seq_idx < mQseqSplitNum; seq_idx++) { run3DKernelDefault(mKernel_split[seq_idx], mGlobalWorkSizeSplit[seq_idx], mLocalWorkSizeSplit[seq_idx], mOpenCLBackend->getOpenCLRuntime()); run3DKernelDefault(mKernel_qk[seq_idx], mGlobalWorkSizeQk[seq_idx], mLocalWorkSizeQk[seq_idx], mOpenCLBackend->getOpenCLRuntime()); run3DKernelDefault(mKernel_softmax[seq_idx], mGlobalWorkSizeSoftMax[seq_idx], mLocalWorkSizeSoftMax[seq_idx], mOpenCLBackend->getOpenCLRuntime()); run3DKernelDefault(mKernel_trans[seq_idx], mGlobalWorkSizeTrans[seq_idx], mLocalWorkSizeTrans[seq_idx], mOpenCLBackend->getOpenCLRuntime()); run3DKernelDefault(mKernel_qkv[seq_idx], mGlobalWorkSizeQkv[seq_idx], mLocalWorkSizeQkv[seq_idx], mOpenCLBackend->getOpenCLRuntime()); run3DKernelDefault(mKernel_clip[seq_idx], mGlobalWorkSizeClip[seq_idx], mLocalWorkSizeClip[seq_idx], mOpenCLBackend->getOpenCLRuntime()); #ifdef DUMP_INTERNAL_LOG { std::ofstream outFile("qk.txt"); std::vector hostPtr_3(16*4096*4096); mOpenCLBackend->getOpenCLRuntime()->commandQueue().enqueueReadBuffer(openCLBuffer(mTempQK.get()), CL_TRUE, 0, 16*4096*4096*4, hostPtr_3.data()); float max_ = -1000000.0; float min_ = 10000000.0; float total = 0.0; for(int i=1; i temp) min_ = temp; } outFile.close(); printf("qk max:%f min:%f avg:%f\n", max_, min_, hostPtr_3[0]+total); } #endif } #endif #ifdef LOG_VERBOSE MNN_PRINT("end SelfAttentionBufExecution onExecute !\n"); #endif return NO_ERROR; } SelfAttentionBufExecution::SelfAttentionBufExecution(const MNN::Op *op, Backend* backend) : CommonExecution(backend, op) { mImpl.reset(new SelfAttentionBufImpl(op, backend)); } SelfAttentionBufExecution::SelfAttentionBufExecution(std::shared_ptr impl, const MNN::Op *op, Backend *backend) : CommonExecution(backend, op), mImpl(impl) {} ErrorCode SelfAttentionBufExecution::onResize(const std::vector& inputs, const std::vector& outputs) { return mImpl->onResize(backend(), inputs, outputs); } ErrorCode SelfAttentionBufExecution::onExecute(const std::vector& inputs, const std::vector& outputs) { return mImpl->onExecute(backend(), inputs, outputs); } bool SelfAttentionBufExecution::onClone(Backend* bn, const Op* op, Execution** dst) { if (nullptr == dst) { return true; } *dst = new SelfAttentionBufExecution(mImpl, op, bn); return true; } class SelfAttentionBufCreator : public OpenCLBackend::Creator { public: virtual Execution *onCreate(const std::vector &inputs, const std::vector &outputs, const MNN::Op *op, Backend *backend) const override { for (int i = 0; i < inputs.size(); ++i) { TensorUtils::setTensorSupportPack(inputs[i], false); } for (int i = 0; i < outputs.size(); ++i) { TensorUtils::setTensorSupportPack(outputs[i], false); } OPENCL_CREATOR_CHECK(new SelfAttentionBufExecution(op, backend)); } }; REGISTER_OPENCL_OP_CREATOR_TRANSFORMER(SelfAttentionBufCreator, OpType_FmhaV2, BUFFER); } // namespace OpenCL } // namespace MNN #endif/* MNN_SUPPORT_TRANSFORMER_FUSE */