diff --git a/modules/dnn/include/opencv2/dnn/all_layers.hpp b/modules/dnn/include/opencv2/dnn/all_layers.hpp index e6fe4c1d4c..905d36f047 100644 --- a/modules/dnn/include/opencv2/dnn/all_layers.hpp +++ b/modules/dnn/include/opencv2/dnn/all_layers.hpp @@ -1800,6 +1800,7 @@ CV__DNN_INLINE_NS_BEGIN bool trans_b; float alpha; float beta; + bool flatten_a; static Ptr create(const LayerParams& params); }; @@ -1808,6 +1809,8 @@ CV__DNN_INLINE_NS_BEGIN public: bool trans_a; bool trans_b; + float alpha; + float beta; static Ptr create(const LayerParams ¶ms); }; diff --git a/modules/dnn/src/graph_fusion_attention.cpp b/modules/dnn/src/graph_fusion_attention.cpp index c9e2e591cc..24f94bb831 100644 --- a/modules/dnn/src/graph_fusion_attention.cpp +++ b/modules/dnn/src/graph_fusion_attention.cpp @@ -119,6 +119,56 @@ struct ModelFusionAttention return isScalarBinOp(prog, idx, NaryEltwiseLayer::OPERATION::DIV, val); } + // True if `arg` is produced by the dynamic scale chain Sqrt<-Cast<-Div(1,.)<-Sqrt<-Cast<-Slice<-Shape; visited ops are appended to `chain_ops`. + bool isRuntimeQKScaleChain(const vector>& prog, Arg arg, + std::set& chain_ops) const + { + const std::vector expected = { + "Sqrt", "Cast2", "NaryEltwise" /*Div*/, "Sqrt", + "Cast2", "Slice2", "Shape" + }; + Arg cur = arg; + for (const std::string& want : expected) { + auto it = producer_.find(cur.idx); + if (it == producer_.end()) return false; + int idx = it->second; + if (idx < 0 || idx >= (int)prog.size() || !prog[idx]) return false; + const Ptr& l = prog[idx]; + if (want == "NaryEltwise") { + NaryEltwiseLayer* elt = dynamic_cast(l.get()); + if (!elt || elt->op != NaryEltwiseLayer::OPERATION::DIV) return false; + + bool one_seen = false; + Arg runtime_in; + bool runtime_seen = false; + for (Arg in : l->inputs) { + if (netimpl->isConstArg(in)) { + Mat t = netimpl->argTensor(in); + if (t.total() != 1) return false; + float v = 0.f; + if (t.type() == CV_32F) v = t.at(0); + else if (t.type() == CV_64F) v = (float)t.at(0); + else return false; + if (std::abs(v - 1.f) > 1e-5f) return false; + one_seen = true; + } else { + runtime_in = in; + runtime_seen = true; + } + } + if (!one_seen || !runtime_seen) return false; + chain_ops.insert(idx); + cur = runtime_in; + } else { + if (l->type != want) return false; + chain_ops.insert(idx); + if (l->inputs.empty()) return false; + cur = l->inputs[0]; + } + } + return true; + } + // Accept Add op with exactly two inputs; identify the non-constant runtime // input (the mask tensor). Returns false if the Add doesn't match. bool isMaskAdd(const vector>& prog, int idx, Arg* out_mask) const @@ -333,11 +383,19 @@ struct ModelFusionAttention { int j = 0; for (int k = 0; k < 3; k++) if (k != k_slot) remaining[j++] = k; } Arg k_tr_out = prog[transpose_idx[k_slot]]->outputs[0]; - int k_next = singleConsumer(k_tr_out); + // Tolerate a Shape consumer alongside the Mul/MatMul: the runtime-scale chain (Shape->Slice->Cast->Sqrt...) branches off the Q/K transpose. + int k_next = findMatchingConsumer(prog, k_tr_out, + [](Layer* L) { + return dynamic_cast(L) != nullptr || + dynamic_cast(L) != nullptr; + }, + &extra_ops); int k_mul_idx = -1; float k_scale_val = 1.f; int qk_matmul_idx = -1; bool vit_style; + bool runtime_qk_scale = false; + std::set qk_scale_chain_ops; if (isScalarMul(prog, k_next, &k_scale_val)) { vit_style = true; k_mul_idx = k_next; @@ -345,6 +403,26 @@ struct ModelFusionAttention } else if (isMatMul(prog, k_next)) { vit_style = false; qk_matmul_idx = k_next; + } else if (k_next >= 0 && k_next < (int)prog.size() && prog[k_next]) { + // K_Transpose -> Mul(K, runtime_scale) where scale = Sqrt(Cast(Div(1,Sqrt(Cast(Slice(Shape(...))))))) and the same scale feeds the Q-side Mul (verified later). + NaryEltwiseLayer* elt = dynamic_cast(prog[k_next].get()); + if (!elt || elt->op != NaryEltwiseLayer::OPERATION::PROD || + prog[k_next]->inputs.size() != 2) { + continue; + } + Arg k_scale_in; + bool tensor_input_seen = false; + for (Arg in : prog[k_next]->inputs) { + if (in.idx == k_tr_out.idx) tensor_input_seen = true; + else k_scale_in = in; + } + if (!tensor_input_seen) continue; + if (!isRuntimeQKScaleChain(prog, k_scale_in, qk_scale_chain_ops)) continue; + + runtime_qk_scale = true; + vit_style = true; + k_mul_idx = k_next; + qk_matmul_idx = singleConsumer(prog[k_mul_idx]->outputs[0]); } else { continue; } @@ -360,8 +438,32 @@ struct ModelFusionAttention Arg q_tr_out = prog[transpose_idx[q_slot]]->outputs[0]; if (vit_style) { - int q_next = singleConsumer(q_tr_out); - if (!isScalarMul(prog, q_next, &q_scale_val)) continue; + int q_next = findMatchingConsumer(prog, q_tr_out, + [](Layer* L) { + return dynamic_cast(L) != nullptr || + dynamic_cast(L) != nullptr; + }, + &extra_ops); + if (isScalarMul(prog, q_next, &q_scale_val)) { + // existing constant-scalar path + } else if (runtime_qk_scale && q_next >= 0 && + q_next < (int)prog.size() && prog[q_next]) { + NaryEltwiseLayer* elt = dynamic_cast(prog[q_next].get()); + if (!elt || elt->op != NaryEltwiseLayer::OPERATION::PROD || + prog[q_next]->inputs.size() != 2) { + continue; + } + Arg q_scale_in; + bool tensor_input_seen = false; + for (Arg in : prog[q_next]->inputs) { + if (in.idx == q_tr_out.idx) tensor_input_seen = true; + else q_scale_in = in; + } + if (!tensor_input_seen) continue; + if (!isRuntimeQKScaleChain(prog, q_scale_in, qk_scale_chain_ops)) continue; + } else { + continue; + } q_mul_idx = q_next; Arg q_mul_out = prog[q_mul_idx]->outputs[0]; int q_dest = singleConsumer(q_mul_out); @@ -463,8 +565,20 @@ struct ModelFusionAttention memcpy(dst + bq.total() + bk.total(), bv.ptr(), bv.total() * sizeof(float)); } - float param_scale = vit_style ? (1.f / (q_scale_val * k_scale_val)) - : post_scale_val; + float param_scale; + if (runtime_qk_scale) { + // Both Q and K were scaled by head_dim^-0.25, so + // (Q'*K'^T) = (1/sqrt(d)) * (Q*K^T). The Attention + // layer's `scale` param is the divisor applied + // before softmax — i.e., sqrt(d). + if (num_heads[0] <= 0 || q_hidden % num_heads[0] != 0) continue; + int head_dim = q_hidden / num_heads[0]; + if (head_dim <= 0) continue; + param_scale = std::sqrt((float)head_dim); + } else { + param_scale = vit_style ? (1.f / (q_scale_val * k_scale_val)) + : post_scale_val; + } LayerParams attn_params; attn_params.name = prog[indices[q_slot]]->name + "_fused_attention"; @@ -504,6 +618,24 @@ struct ModelFusionAttention if (qk_div_idx >= 0) to_remove.insert(qk_div_idx); if (qk_add_idx >= 0) to_remove.insert(qk_add_idx); for (int op : extra_ops) to_remove.insert(op); + + // Drop chain ops (Sqrt/Cast/Div/Slice/Shape) only if every consumer is also being removed — by this fusion or another chain op already in the set; the upstream is shared between Q and K branches so the chain is treated as a unit. + for (int op_idx : qk_scale_chain_ops) { + if (op_idx < 0 || op_idx >= (int)prog.size() || !prog[op_idx]) continue; + bool all_consumers_removed = true; + for (Arg out : prog[op_idx]->outputs) { + auto cit = consumers_.find(out.idx); + if (cit == consumers_.end()) continue; + for (int c : cit->second) { + if (to_remove.count(c) || qk_scale_chain_ops.count(c)) continue; + all_consumers_removed = false; + break; + } + if (!all_consumers_removed) break; + } + if (all_consumers_removed) to_remove.insert(op_idx); + } + for (int op : to_remove) removed_ops.insert(op); diff --git a/modules/dnn/src/graph_fusion_matmul_to_gemm.cpp b/modules/dnn/src/graph_fusion_matmul_to_gemm.cpp new file mode 100644 index 0000000000..b3f8ee4096 --- /dev/null +++ b/modules/dnn/src/graph_fusion_matmul_to_gemm.cpp @@ -0,0 +1,148 @@ +// This file is part of OpenCV project. +// It is subject to the license terms in the LICENSE file found in the top-level directory +// of this distribution and at http://opencv.org/license.html. + +// Rewrites a MatMul whose B (and optionally bias) are constant blobs into a +// Gemm with `flatten_a=false`. The Gemm path pre-packs B in MLAS layout at +// finalize() and dispatches each forward through MlasGemmPacked, which is +// notably faster than the OpenCV-internal packed AVX2 kernel that the MatMul +// constant-B path otherwise hits. +// +// Pattern (post-importer; bias has already been absorbed into MatMul.blobs[1] +// by BiasedMatmulSubgraph when it was a const Add): +// A(any rank) -> MatMul[const B (2D), optional const bias [N] / scalar] +// +// Replacement: +// A -> Gemm[transA=false, transB=mm.trans_b, alpha=mm.alpha, beta=mm.beta, +// constB=true, const_C=have_bias, flatten_a=false, +// blobs={B, bias?}] +// +// Restrictions (skip otherwise): +// - MatMul.trans_a must be false (ND-A flattening assumes row-major K is the +// last axis) +// - blobs.size() in {1, 2} and blobs[0] (B) is 2D float32 +// - if a bias is present, it must be scalar (total==1) or 1D length-N. The +// Gemm flatten_a=false bias path tiles a per-row pattern, so 2D / per-row +// biases like [M, N] aren't supported by this rewriter (they would change +// value across the flattened rows). + +#include "precomp.hpp" +#include "net_impl.hpp" + +namespace cv { namespace dnn { +CV__DNN_INLINE_NS_BEGIN + +using std::vector; +using std::string; + +struct ModelFusionMatMulToGemm +{ + explicit ModelFusionMatMulToGemm(Net::Impl* netimpl_) : netimpl(netimpl_) {} + + void fuse() { fuseGraph(netimpl->mainGraph); } + + bool fuseGraph(Ptr& graph) + { + const vector>& prog = graph->prog(); + size_t nops = prog.size(); + bool modified = false; + + for (size_t i = 0; i < nops; i++) { + if (!prog[i]) continue; + vector>* subgraphs = prog[i]->subgraphs(); + if (subgraphs) { + for (Ptr& g : *subgraphs) + if (fuseGraph(g)) modified = true; + } + } + + vector> newprog = prog; + bool changed = false; + + for (size_t i = 0; i < nops; i++) { + const Ptr& layer = newprog[i]; + if (!layer) continue; + + MatMulLayer* mm = dynamic_cast(layer.get()); + if (!mm) continue; + // MatMulInt8Layer carries quantization state that this rewriter + // doesn't know how to translate to Gemm. + if (dynamic_cast(layer.get())) continue; + + // Constant B lives in blobs[0] post-parse; if blobs is empty the + // MatMul has a runtime B and isn't a candidate. + if (layer->blobs.empty() || layer->blobs.size() > 2) continue; + + if (mm->trans_a) continue; // ND-A flatten assumes K = last axis + + const Mat& B = layer->blobs[0]; + if (B.dims != 2 || B.type() != CV_32F) continue; + + // Single runtime input expected (the const B was absorbed into blobs). + if (layer->inputs.size() != 1) continue; + + int N = mm->trans_b ? B.size[0] : B.size[1]; + int K = mm->trans_b ? B.size[1] : B.size[0]; + (void)K; + if (N <= 0) continue; + + const bool have_bias = layer->blobs.size() == 2; + if (have_bias) { + const Mat& bias = layer->blobs[1]; + if (bias.type() != CV_32F) continue; + int total = (int)bias.total(); + if (total != 1 && total != N) continue; + } + + // Build the Gemm replacement. + LayerParams gp; + gp.name = layer->name; // keep the name for profiling continuity + gp.type = "Gemm"; + gp.set("transA", false); + gp.set("transB", mm->trans_b); + gp.set("alpha", mm->alpha); + gp.set("beta", mm->beta); + gp.set("constB", true); + gp.set("have_bias", have_bias); + gp.set("const_C", have_bias); + // flatten_a=false: keep A's leading dims so downstream consumers + // see the same shape they did when the producer was a MatMul. + gp.set("flatten_a", false); + // For the new GemmLayerImpl flatten_a=false bias path the bias is + // restricted to scalar or [N]; both of those map to real_ndims_C + // <= 1, but we don't actually consult it in that path — set it + // anyway so the Ngraph/CANN backends still get a sensible value. + if (have_bias) { + int total = (int)layer->blobs[1].total(); + gp.set("real_ndims_C", total == 1 ? 0 : 1); + } + + gp.blobs.push_back(B); + if (have_bias) gp.blobs.push_back(layer->blobs[1]); + + Ptr gemm = LayerFactory::createLayerInstance("Gemm", gp); + if (!gemm) continue; + gemm->inputs = layer->inputs; + gemm->outputs = layer->outputs; + gemm->netimpl = netimpl; + + newprog[i] = gemm; + changed = true; + modified = true; + } + + if (changed) graph->setProg(newprog); + return modified; + } + + Net::Impl* netimpl; +}; + +void Net::Impl::fuseMatMulConstBToGemm() +{ + ModelFusionMatMulToGemm pass(this); + pass.fuse(); +} + +CV__DNN_INLINE_NS_END +}} diff --git a/modules/dnn/src/graph_fusion_shared_gemm.cpp b/modules/dnn/src/graph_fusion_shared_gemm.cpp index 4f0ab67947..643adfd23d 100644 --- a/modules/dnn/src/graph_fusion_shared_gemm.cpp +++ b/modules/dnn/src/graph_fusion_shared_gemm.cpp @@ -73,6 +73,7 @@ struct ModelFusionSharedGemm float alpha = 1.f; float beta = 1.f; bool has_bias = false; + bool flatten_a = true; }; bool inspectGemm(const vector>& prog, int idx, GemmInfo& info) const @@ -116,6 +117,7 @@ struct ModelFusionSharedGemm info.alpha = alpha; info.beta = beta; info.has_bias = (l->blobs.size() == 2); + info.flatten_a = g->flatten_a; return true; } @@ -154,6 +156,14 @@ struct ModelFusionSharedGemm if (info.has_bias != all_have_bias) { uniform = false; break; } if (!uniform) continue; + // Require uniform flatten_a too — the fused Gemm has a single + // value, and mixing 2D-output and ND-output downstream consumers + // would need different post-fusion shape handling. + bool all_flatten_a = infos[0].flatten_a; + for (auto& info : infos) + if (info.flatten_a != all_flatten_a) { uniform = false; break; } + if (!uniform) continue; + int K = infos[0].K; int total_N = 0; for (auto& info : infos) total_N += info.N; @@ -203,6 +213,15 @@ struct ModelFusionSharedGemm fp.set("transB", false); fp.set("alpha", 1.f); fp.set("beta", 1.f); + fp.set("flatten_a", all_flatten_a); + // Mirror the constB / const_C / have_bias signalling that the + // GemmLayerImpl's getOpMode() reads from LayerParams. We always + // ship the weights as a constant blob, and (when biased) the bias + // too — so this fused Gemm has only one runtime input. + fp.set("constB", true); + fp.set("have_bias", all_have_bias); + fp.set("const_C", all_have_bias); + if (all_have_bias) fp.set("real_ndims_C", 1); fp.blobs.push_back(W_concat); if (all_have_bias) fp.blobs.push_back(b_concat); diff --git a/modules/dnn/src/layers/attention_layer.cpp b/modules/dnn/src/layers/attention_layer.cpp index df70dfbdad..5a4800a8dd 100644 --- a/modules/dnn/src/layers/attention_layer.cpp +++ b/modules/dnn/src/layers/attention_layer.cpp @@ -69,11 +69,11 @@ static void rotationKernel( } } }; - // This will spin up threads and run fn over [0, seq_len) parallel_for_(cv::Range(0, int(seq_len)), fn, nstripes); } +// Precomputes RoPE sin/cos table of shape [seq_len, d] (https://arxiv.org/pdf/2104.09864). static void precompRotationTable(float *data, size_t seq_len, size_t d) { @@ -299,28 +299,25 @@ class AttentionLayerImpl CV_FINAL : public AttentionLayer { if (do_rotary) { - // precompute sin/cos table auto &rope_table = internals.back(); auto *rope_table_data = rope_table.ptr(); - // currently, support rotary embeddings only if q and k head sizes are equal + // RoPE currently requires q and k head sizes to match. CV_Assert(qkv_head_sizes[0] == qkv_head_sizes[1]); precompRotationTable(rope_table_data, seq_len, qkv_head_sizes[0]); } - // Compute Q/K/V auto &gemm_buffer = internals[0]; auto *Q = gemm_buffer.ptr(); auto *K = Q + batch_size * seq_len * qkv_hidden_sizes[0]; auto *V = K + batch_size * seq_len * qkv_hidden_sizes[1]; - float *QKV[3] = {Q, K, V}; // Q, K, V: [B, N, S, H] + float *QKV[3] = {Q, K, V}; // [B, N, S, H] per tensor { const auto &input = inputs[0]; const auto &bias = blobs.empty() ? inputs[2] : blobs.back(); const auto *input_data = input.ptr(); const auto *bias_data = bias.ptr(); - // If rotary is false, evaluates to internals[2], which is the output_buffer - // but this is not dramatic, because in case rotary is false, the table is not used + // When do_rotary is false, internals.back() aliases output_buffer; harmless because rope_table is then unused. const auto &rope_table = internals.back(); opt.multi_thread = false; @@ -337,7 +334,7 @@ class AttentionLayerImpl CV_FINAL : public AttentionLayer { int bias_offset = qkv_index * qkv_hidden_sizes[0] + head_index * head_size; int dst_offset = (batch_index * num_heads + head_index) * (seq_len * head_size); - // broadcast bias ([NH] -> [BN, SH]) and make copy to dst + // Broadcast bias [NH] -> [BN, SH] into dst before gemm so beta=1 adds it implicitly. const auto *bias_data_src = bias_data + bias_offset; auto *dst_data = dst + dst_offset; for (size_t seq_len_idx = 0; seq_len_idx < seq_len; seq_len_idx++) { @@ -346,13 +343,13 @@ class AttentionLayerImpl CV_FINAL : public AttentionLayer { } auto *packed_weight = packed_weights[qkv_index] + packed_weights_size[qkv_index] * head_index; - // single-thread gemm kernel + // Single-threaded gemm; outer parallel_for_ already partitions over (batch, head, qkv). fastGemm(false, seq_len, head_size, input_hidden_size, 1.f, input_data + input_offset, input_hidden_size, packed_weight, 1.f, dst + dst_offset, head_size, opt); if(qkv_index < 2 && do_rotary) { - // rotate on the fly + // Apply RoPE to Q/K in place (V is left untouched). const auto *rope_table_data = rope_table.ptr(); rotationKernel( dst + dst_offset, @@ -369,7 +366,7 @@ class AttentionLayerImpl CV_FINAL : public AttentionLayer { parallel_for_(Range(0, loops), fn, nstripes); } - // Compute Softmax(scale * MatMul(Q, K)) + // attention_prob = softmax(scale * Q @ K^T + mask) auto &attention_prob = internals[1]; { auto *output = attention_prob.ptr(); @@ -379,23 +376,22 @@ class AttentionLayerImpl CV_FINAL : public AttentionLayer { auto qk_head_size = qkv_head_sizes[0]; auto qk_inner_size = seq_len * qk_head_size; - // Compute scale * matmul(Q, K) - opt.multi_thread = false; - parallel_for_(Range(0, loops), [&] (const Range r) { - for (int i = r.start; i < r.end; i++) { - const int output_offset = i * seq_len_square; + // One batched fastGemmBatch call over all (batch, head) pairs — wrapping per-head fastGemm in parallel_for_ caused nested parallelism (inner MLAS gemm issued its own parallel_for_, serializing on most threads). + std::vector qk_a_offs(loops), qk_b_offs(loops), qk_c_offs(loops); + for (int i = 0; i < (int)loops; i++) { + qk_a_offs[i] = (size_t)i * qk_inner_size; + qk_b_offs[i] = (size_t)i * qk_inner_size; + qk_c_offs[i] = (size_t)i * seq_len_square; + } + opt.multi_thread = true; + // ldb0=1, ldb1=qk_head_size signals trans_b for K (stored as [seq_len, qk_head_size]) to fastGemmBatch. + fastGemmBatch(loops, qk_a_offs.data(), qk_b_offs.data(), qk_c_offs.data(), + (int)seq_len, (int)seq_len, (int)qk_head_size, + scale, Q, (int)qk_head_size, 1, + K, 1, (int)qk_head_size, 0.f, + output, (int)seq_len, opt); - const auto *q = Q + qk_inner_size * i, *k = K + qk_inner_size * i; - fastGemm(false, true, seq_len, qk_head_size, seq_len, qk_head_size, - scale, q, qk_head_size, 1, - k, qk_head_size, 1, 0.f, - output + output_offset, seq_len, opt); - } - }, loops * seq_len * qk_head_size * seq_len * (1 / 1024.0)); - - // Additive attention mask applied before softmax. Supports BERT-style mask - // shapes [B, 1, 1, S] (per-batch) and [1, 1, 1, S] (global broadcast). - // Only reached when the fusion pass attached an external mask input. + // Additive mask broadcast-aligned right to [B,H,S,S] (size-1 dims broadcast); only present when the fusion pass attached an external mask input. int num_non_blob_inputs = (int)inputs.size(); bool has_mask = (!blobs.empty() && num_non_blob_inputs >= 2) || (blobs.empty() && num_non_blob_inputs >= 4); @@ -403,36 +399,72 @@ class AttentionLayerImpl CV_FINAL : public AttentionLayer { const Mat &mask_mat = blobs.empty() ? inputs[3] : inputs[1]; CV_CheckTypeEQ(mask_mat.type(), CV_32F, "DNN/Attention: mask must be float32"); - const int mask_last_dim = shape(mask_mat).back(); - CV_CheckEQ(mask_last_dim, (int)seq_len, - "DNN/Attention: mask last dim must equal seq_len"); - const size_t mask_total = mask_mat.total(); - int mask_batch_stride; - if (mask_total == batch_size * seq_len) mask_batch_stride = (int)seq_len; - else if (mask_total == seq_len) mask_batch_stride = 0; - else CV_Error(Error::StsNotImplemented, - "DNN/Attention: unsupported mask shape"); + + const auto mask_shape = shape(mask_mat); + const int mask_ndim = mask_shape.dims; + auto fmt_shape = [&]() { + std::ostringstream ss; + for (int d = 0; d < mask_ndim; d++) { if (d) ss << ","; ss << mask_shape[d]; } + return ss.str(); + }; + // Right-align mask dims to (B, H, Q=S, K=S); missing leading dims default to 1. + auto get_dim = [&](int t) -> int { + int md = mask_ndim - 4 + t; + return (md >= 0) ? mask_shape[md] : 1; + }; + // Extra leading dims beyond rank 4 must all be 1 (no broadcasting beyond [B,H,S,S]). + for (int d = 0; d < mask_ndim - 4; d++) { + if (mask_shape[d] != 1) + CV_Error(Error::StsNotImplemented, + cv::format("DNN/Attention: unsupported mask shape [%s] (B=%d, H=%d, S=%d)", + fmt_shape().c_str(), (int)batch_size, (int)num_heads, (int)seq_len)); + } + const int dim_b = get_dim(0); + const int dim_h = get_dim(1); + const int dim_q = get_dim(2); + const int dim_k = get_dim(3); + auto check_bcast = [&](int md, int td, const char* axis) { + if (md != 1 && md != td) + CV_Error(Error::StsNotImplemented, + cv::format("DNN/Attention: mask shape [%s] not broadcastable to [%d,%d,%d,%d] (axis %s)", + fmt_shape().c_str(), (int)batch_size, (int)num_heads, + (int)seq_len, (int)seq_len, axis)); + }; + check_bcast(dim_b, (int)batch_size, "batch"); + check_bcast(dim_h, (int)num_heads, "head"); + check_bcast(dim_q, (int)seq_len, "query"); + check_bcast(dim_k, (int)seq_len, "key"); + + const size_t mask_b_stride = (dim_b == 1) ? 0 : (size_t)dim_h * dim_q * dim_k; + const size_t mask_h_stride = (dim_h == 1) ? 0 : (size_t)dim_q * dim_k; + const size_t mask_q_stride = (dim_q == 1) ? 0 : (size_t)dim_k; + const bool mask_k_bcast = (dim_k == 1); const float *mask_data = mask_mat.ptr(); parallel_for_(Range(0, (int)loops), [&](const Range &r) { for (int i = r.start; i < r.end; i++) { const int b = i / (int)num_heads; - const float *m = mask_data + b * mask_batch_stride; + const int h = i % (int)num_heads; + const float *m_bh = mask_data + b * mask_b_stride + h * mask_h_stride; float *prob = output + i * seq_len_square; for (size_t row = 0; row < seq_len; row++) { + const float *m = m_bh + row * mask_q_stride; float *p = prob + row * seq_len; - for (size_t j = 0; j < seq_len; j++) - p[j] += m[j]; + if (mask_k_bcast) { + float v = m[0]; + for (size_t j = 0; j < seq_len; j++) p[j] += v; + } else { + for (size_t j = 0; j < seq_len; j++) p[j] += m[j]; + } } } }, loops * seq_len * (1 / 1024.0)); } - // Compute softmax on the last dimension softmax(attention_prob, attention_prob, shape(attention_prob).size() - 1); } - // Compute MatMul(attention_prob, V) + // output = attention_prob @ V, then transpose back to [B, S, H*D] auto &output_buffer = internals[2]; { auto *output = outputs[0].ptr(); @@ -444,29 +476,35 @@ class AttentionLayerImpl CV_FINAL : public AttentionLayer { auto v_head_size = qkv_head_sizes[2]; auto v_inner_size = seq_len * v_head_size; - opt.multi_thread = false; - parallel_for_(Range(0, loops), [&] (const Range &r) { + // Batched fastGemmBatch over (batch, head) — same nested-parallelism rationale as QK^T above. + std::vector av_a_offs(loops), av_b_offs(loops), av_c_offs(loops); + for (int i = 0; i < (int)loops; i++) { + av_a_offs[i] = (size_t)i * prob_inner_size; + av_b_offs[i] = (size_t)i * v_inner_size; + av_c_offs[i] = (size_t)i * v_inner_size; + } + opt.multi_thread = true; + fastGemmBatch(loops, av_a_offs.data(), av_b_offs.data(), av_c_offs.data(), + (int)seq_len, (int)v_head_size, (int)seq_len, + 1.f, prob, (int)seq_len, 1, + V, (int)v_head_size, 1, 0.f, + output_buff, (int)v_head_size, opt); + + // Transpose [B*H, S, D] -> [B, S, H*D] in place via per-(batch, head) memcpy strips. + parallel_for_(Range(0, (int)loops), [&] (const Range &r) { for (int i = r.start; i < r.end; i++) { - const int output_offset = i * v_inner_size; - - const auto *p = prob + i * prob_inner_size, *v = V + i * v_inner_size; - fastGemm(false, false, seq_len, seq_len, seq_len, v_head_size, - 1.f, p, seq_len, 1, - v, v_head_size, 1, 0.f, - output_buff + output_offset, v_head_size, opt); - - // tranpose on the fly + const int output_offset = i * (int)v_inner_size; const int batch_index = static_cast(i / num_heads); - const int head_index = static_cast(i % num_heads); - auto *src = output_buff + output_offset; - auto *dst = output + (batch_index * seq_len * num_heads + head_index) * v_head_size; - for (int j = 0; j < seq_len; j++) { + const int head_index = static_cast(i % num_heads); + const float *src = output_buff + output_offset; + float *dst = output + (batch_index * (int)seq_len * (int)num_heads + head_index) * (int)v_head_size; + for (int j = 0; j < (int)seq_len; j++) { std::memcpy(dst, src, v_head_size * sizeof(float)); src += v_head_size; dst += qkv_hidden_sizes[2]; } } - }, loops * seq_len * seq_len * v_head_size * (1 / 1024.0)); + }, loops * seq_len * v_head_size * (1 / 1024.0)); } } diff --git a/modules/dnn/src/layers/concat2_layer.cpp b/modules/dnn/src/layers/concat2_layer.cpp index 8e75529a2a..9878f9bf89 100644 --- a/modules/dnn/src/layers/concat2_layer.cpp +++ b/modules/dnn/src/layers/concat2_layer.cpp @@ -157,6 +157,33 @@ public: internals.clear(); } + int getLayouts(const std::vector& actualInputs, + std::vector& desiredInputs, + const int requiredOutputs, + std::vector& outputs) const CV_OVERRIDE + { + auto* netimpl_ = getNetImpl(this); + DataLayout defaultLayout = netimpl_->originalLayout; + const size_t ninputs = actualInputs.size(); + desiredInputs = actualInputs; + outputs.assign(requiredOutputs, DATA_LAYOUT_UNKNOWN); + + bool allBlock = ninputs > 0; + for (size_t i = 0; i < ninputs; ++i) + if (actualInputs[i] != DATA_LAYOUT_BLOCK) { allBlock = false; break; } + + // BLOCK layout on the channel axis would expose inner-block padding as real channels; let TransformLayout repack instead. + const bool canKeepBlock = allBlock && axis >= 0 && axis != 1; + + if (canKeepBlock) { + outputs.assign(requiredOutputs, DATA_LAYOUT_BLOCK); + } else { + for (size_t i = 0; i < ninputs; ++i) + if (actualInputs[i] == DATA_LAYOUT_BLOCK) desiredInputs[i] = defaultLayout; + } + return outputs[0] == DATA_LAYOUT_BLOCK ? netimpl_->defaultC0 : 0; + } + void finalize(InputArrayOfArrays, OutputArrayOfArrays outputs_arr) CV_OVERRIDE { } diff --git a/modules/dnn/src/layers/conv2_layer.cpp b/modules/dnn/src/layers/conv2_layer.cpp index d6d46a1d84..e90b4390e7 100644 --- a/modules/dnn/src/layers/conv2_layer.cpp +++ b/modules/dnn/src/layers/conv2_layer.cpp @@ -287,7 +287,7 @@ public: desiredInputs[0] = DATA_LAYOUT_BLOCK; for (size_t i = 1; i < ninputs; i++) desiredInputs[i] = DATA_LAYOUT_UNKNOWN; - if (addResidual && ninputs >= 2) + if (addResidual && ninputs > 1) desiredInputs[ninputs - 1] = DATA_LAYOUT_BLOCK; outputs.assign(requiredOutputs, DATA_LAYOUT_BLOCK); return getNetImpl(this)->defaultC0; diff --git a/modules/dnn/src/layers/cpu_kernels/fast_gemm.cpp b/modules/dnn/src/layers/cpu_kernels/fast_gemm.cpp index e7e5605cbb..9e87590d94 100644 --- a/modules/dnn/src/layers/cpu_kernels/fast_gemm.cpp +++ b/modules/dnn/src/layers/cpu_kernels/fast_gemm.cpp @@ -573,11 +573,23 @@ void fastGemmBatch(size_t batch, const size_t *A_offsets, const size_t *B_offset // Below ~10K MACs (M*N*K) the blocked SIMD kernel's pack/tile overhead outweighs its speedup, so route tiny problems through the unblocked thin path. if (ldb1 == 1 && (fastGemmThinEligible(M, N, K) || (uint64_t)M * N * K <= 10000)) { const size_t esz = sizeof(float); - for (size_t i = 0; i < batch; i++) { - fast_gemm_thin(alpha, beta, M, N, K, - a + A_offsets[i] * esz, lda0, lda1, - b + B_offsets[i] * esz, ldb0, - c + C_offsets[i] * esz, ldc, opt.multi_thread); + // Parallelise over batch with single-thread inner thin gemm + if (opt.multi_thread && batch > 1) { + parallel_for_(Range(0, (int)batch), [&](const Range& r) { + for (int i = r.start; i < r.end; i++) { + fast_gemm_thin(alpha, beta, M, N, K, + a + A_offsets[i] * esz, lda0, lda1, + b + B_offsets[i] * esz, ldb0, + c + C_offsets[i] * esz, ldc, false); + } + }, (double)batch * M * N * K * (1.0 / 1024.0)); + } else { + for (size_t i = 0; i < batch; i++) { + fast_gemm_thin(alpha, beta, M, N, K, + a + A_offsets[i] * esz, lda0, lda1, + b + B_offsets[i] * esz, ldb0, + c + C_offsets[i] * esz, ldc, opt.multi_thread); + } } return; } @@ -675,11 +687,22 @@ void fastGemmBatch(size_t batch, if (ldb1 == 1 && (fastGemmThinEligible(M, N, K) || (uint64_t)M * N * K <= 10000)) { const size_t esz = sizeof(float); - for (size_t i = 0; i < batch; i++) { - fast_gemm_thin(alpha, beta, M, N, K, - a + A_offsets[i] * esz, lda0, lda1, - b + B_offsets[i] * esz, ldb0, - c + C_offsets[i] * esz, ldc, opt.multi_thread); + if (opt.multi_thread && batch > 1) { + parallel_for_(Range(0, (int)batch), [&](const Range& r) { + for (int i = r.start; i < r.end; i++) { + fast_gemm_thin(alpha, beta, M, N, K, + a + A_offsets[i] * esz, lda0, lda1, + b + B_offsets[i] * esz, ldb0, + c + C_offsets[i] * esz, ldc, false); + } + }, (double)batch * M * N * K * (1.0 / 1024.0)); + } else { + for (size_t i = 0; i < batch; i++) { + fast_gemm_thin(alpha, beta, M, N, K, + a + A_offsets[i] * esz, lda0, lda1, + b + B_offsets[i] * esz, ldb0, + c + C_offsets[i] * esz, ldc, opt.multi_thread); + } } return; } diff --git a/modules/dnn/src/layers/cpu_kernels/fast_gemm_kernels.simd.hpp b/modules/dnn/src/layers/cpu_kernels/fast_gemm_kernels.simd.hpp index b9ce5a9113..1f9eac249f 100644 --- a/modules/dnn/src/layers/cpu_kernels/fast_gemm_kernels.simd.hpp +++ b/modules/dnn/src/layers/cpu_kernels/fast_gemm_kernels.simd.hpp @@ -666,8 +666,7 @@ void fastGemmKernel(size_t M, size_t N, size_t K, }; if (multi_thread) { - int cost_per_thread = static_cast((K / KC) * (MC / GEMM_MR) * (NC / GEMM_NR)); - double nstripes = (size_t)total_tiles * cost_per_thread * (1 / 1024.0); + double nstripes = (double)M * N * K * (1.0 / 1024.0); parallel_for_(Range(0, total_tiles), fn, nstripes); } else { fn(Range(0, total_tiles)); @@ -750,8 +749,7 @@ void fastGemmKernel(size_t M, size_t N, size_t K, }; if (multi_thread) { - int cost_per_thread = static_cast((K / KC) * (MC / GEMM_MR) * (NC / GEMM_NR)); - double nstripes = (size_t)total_tiles * cost_per_thread * (1 / 1024.0); + double nstripes = (double)M * N * K * (1.0 / 1024.0); parallel_for_(Range(0, total_tiles), fn, nstripes); } else { fn(Range(0, total_tiles)); @@ -846,8 +844,7 @@ void fastGemmBatchKernel(size_t batch, const size_t *A_offsets, const size_t *B_ }; int total = batch * total_tiles; - int cost_per_thread = static_cast((K / KC) * (MC / GEMM_MR) * (NC / GEMM_NR)); - double nstripes = (size_t)total * cost_per_thread * (1 / 1024.0); + double nstripes = (double)batch * M * N * K * (1.0 / 1024.0); parallel_for_(Range(0, total), fn, nstripes); } @@ -932,8 +929,7 @@ void fastGemmBatchKernel(size_t batch, const size_t *A_offsets, const size_t *B_ }; int total = batch * total_tiles; - int cost_per_thread = static_cast((K / KC) * (MC / GEMM_MR) * (NC / GEMM_NR)); - double nstripes = (size_t)total * cost_per_thread * (1 / 1024.0); + double nstripes = (double)batch * M * N * K * (1.0 / 1024.0); parallel_for_(Range(0, total), fn, nstripes); } diff --git a/modules/dnn/src/layers/flatten_layer.cpp b/modules/dnn/src/layers/flatten_layer.cpp index 8ed9d13235..e628bc3b42 100644 --- a/modules/dnn/src/layers/flatten_layer.cpp +++ b/modules/dnn/src/layers/flatten_layer.cpp @@ -89,6 +89,8 @@ public: bool isDataShuffling() const CV_OVERRIDE { return true; } + virtual bool alwaysSupportInplace() const CV_OVERRIDE { return true; } + bool getMemoryShapes(const std::vector &inputs, const int requiredOutputs, std::vector &outputs, @@ -209,18 +211,11 @@ public: outputs_arr.isUMatVector(), forward_ocl(inputs_arr, outputs_arr, internals_arr)) - std::vector inputs, outputs; - inputs_arr.getMatVector(inputs); - outputs_arr.getMatVector(outputs); - - for (size_t i = 0; i < inputs.size(); i++) - { - MatShape outShape = shape(outputs[i]); - if (inputs[i].data != outputs[i].data) - { - inputs[i].reshape(1, (int)outShape.size(), &outShape[0]).copyTo(outputs[i]); - } - } + std::vector outs; + outputs_arr.getMatVector(outs); + CV_Assert(!outs.empty()); + const MatShape outShape = outs[0].shape(); + reshapeAndCopyFirst(inputs_arr, outputs_arr, outShape); } #ifdef HAVE_CANN diff --git a/modules/dnn/src/layers/gemm_layer.cpp b/modules/dnn/src/layers/gemm_layer.cpp index ceb12f0c8a..0b60165220 100644 --- a/modules/dnn/src/layers/gemm_layer.cpp +++ b/modules/dnn/src/layers/gemm_layer.cpp @@ -58,6 +58,7 @@ public: trans_b = params.get("transB", false); alpha = params.get("alpha", 1.0f); beta = params.get("beta", 1.0f); + flatten_a = params.get("flatten_a", true); // The params are not part of ONNX, but set by old ONNX parser const_B = params.get("constB", false); @@ -164,9 +165,19 @@ public: } } - int batches = std::accumulate(shape_A.begin(), shape_A.end() - 2, 1, std::multiplies()); - MatShape shape_y{M * batches, N}; - outputs.assign(1, shape_y); + if (flatten_a) { + int batches = std::accumulate(shape_A.begin(), shape_A.end() - 2, 1, std::multiplies()); + MatShape shape_y{M * batches, N}; + outputs.assign(1, shape_y); + } else { + // Preserve A's leading dims; only the trailing axis changes from K to N. + // (trans_a is rejected upstream for this mode, so M corresponds to + // shape_A[-2] and we just rewrite the last axis.) + CV_CheckFalse(trans_a, "DNN/Gemm: flatten_a=false requires trans_a=false"); + MatShape shape_y = shape_A; + shape_y[shape_y.size() - 1] = N; + outputs.assign(1, shape_y); + } return false; } @@ -245,10 +256,32 @@ public: // pack B if it is const if (constB(mode)) { fastGemmPackB(blobs[0], packed_B, trans_b, opt); + + // Pre-pack B in the "thin" layout when the gemm shape has a + // small leading dim (M <= FAST_GEMM_THIN_MAX_M). + thin_packed_B.clear(); + if (!trans_a && blobs[0].type() == CV_32F) { + std::vector outputs; + outputs_arr.getMatVector(outputs); + if (!outputs.empty()) { + const auto &Y = outputs[0]; + const auto shape_Y = shape(Y); + const int N = shape_Y.back(); + const int K = trans_b ? blobs[0].size[1] : blobs[0].size[0]; + const int rows_thin = flatten_a ? shape_Y[shape_Y.size() - 2] + : (int)(Y.total() / (size_t)N); + if (fastGemmThinEligible(rows_thin, N, K)) { + thin_packed_B.resize(fastGemmThinPackBSize(N, K)); + const size_t ldb_K = trans_b ? 1 : N; + const size_t ldb_N = trans_b ? K : 1; + fastGemmThinPackB(N, K, blobs[0].ptr(), + ldb_K, ldb_N, thin_packed_B.data()); + } + } + } } - // also pre-broadcast bias - if (constC(mode)) { + if (constC(mode) && flatten_a) { const auto &C = blobs.back(); std::vector outputs; @@ -290,15 +323,37 @@ public: int M = shape_Y[dims_Y - 2], N = shape_Y[dims_Y - 1]; int K = trans_a ? ma : na; + // In flatten_a=false mode the output keeps A's leading dims, so the + // GEMM row count spans those dims as well: rows = total(Y)/N. + const int rows = flatten_a ? M : (int)(Y.total() / (size_t)N); + // broadcast C and copy C to output if (constC(mode) || inputs.size() >= 3) { - if (!constC(mode) || broadcast_C.empty()) { - broadcastCWtihBeta(M, N, (inputs.size() >= 3 ? inputs.back() : blobs.back())); - } - int step = M * N; - CV_CheckEQ(broadcast_C.size(), static_cast(step), "DNN/Gemm: C is not broadcast properly"); float *ptr_y = Y.ptr(); - std::memcpy(ptr_y, broadcast_C.data(), step * sizeof(float)); + if (flatten_a) { + if (!constC(mode) || broadcast_C.empty()) { + broadcastCWtihBeta(M, N, (inputs.size() >= 3 ? inputs.back() : blobs.back())); + } + int step = M * N; + CV_CheckEQ(broadcast_C.size(), static_cast(step), "DNN/Gemm: C is not broadcast properly"); + std::memcpy(ptr_y, broadcast_C.data(), step * sizeof(float)); + } else { + // ND output: tile the (1D / scalar) bias across all `rows` + // rows, scaled by beta. The rewriter restricts bias to scalar + // or 1D length-N; assert here. + const Mat& C = (inputs.size() >= 3) ? inputs.back() : blobs.back(); + const float* c = C.ptr(); + if (C.total() == 1) { + float val = beta * (*c); + std::fill_n(ptr_y, (size_t)rows * (size_t)N, val); + } else { + CV_CheckEQ((int)C.total(), N, "DNN/Gemm: bias must be scalar or length-N in flatten_a=false mode"); + for (int j = 0; j < N; j++) ptr_y[j] = beta * c[j]; + for (int i = 1; i < rows; i++) { + std::memcpy(ptr_y + (size_t)i * N, ptr_y, (size_t)N * sizeof(float)); + } + } + } } else { // initialization float *ptr_y = Y.ptr(); size_t total = Y.total(); @@ -307,7 +362,12 @@ public: if (constB(mode)) { CV_CheckGT(packed_B.size(), static_cast(0), "DNN/Gemm: constant B is not pre-packed"); - fastGemm(trans_a, M, N, K, alpha, A.ptr(), na, packed_B.data(), 1.f, Y.ptr(), N, opt); + if (!thin_packed_B.empty()) { + fastGemmThin(rows, N, K, alpha, A.ptr(), na, 1, + thin_packed_B.data(), 1.f, Y.ptr(), N, opt.multi_thread); + } else { + fastGemm(trans_a, rows, N, K, alpha, A.ptr(), na, packed_B.data(), 1.f, Y.ptr(), N, opt); + } } else { fastGemmBatch(trans_a, trans_b, alpha, A, inputs[1], 1.f, Y, opt); } @@ -470,6 +530,7 @@ private: bool const_C; bool have_bias; std::vector packed_B; + std::vector thin_packed_B; std::vector broadcast_C; int real_ndims_C; FastGemmOpt opt; diff --git a/modules/dnn/src/layers/matmul_layer.cpp b/modules/dnn/src/layers/matmul_layer.cpp index e0c70ffc65..7f84cbef06 100644 --- a/modules/dnn/src/layers/matmul_layer.cpp +++ b/modules/dnn/src/layers/matmul_layer.cpp @@ -503,9 +503,6 @@ class MatMulLayerImpl CV_FINAL : public MatMulLayer { #endif // HAVE_CANN private: - float alpha; - float beta; - int real_ndims_C; std::vector packed_input_B; diff --git a/modules/dnn/src/layers/slice2_layer.cpp b/modules/dnn/src/layers/slice2_layer.cpp index ce12fb1631..b135068509 100644 --- a/modules/dnn/src/layers/slice2_layer.cpp +++ b/modules/dnn/src/layers/slice2_layer.cpp @@ -130,6 +130,42 @@ public: bool isDataShuffling() const CV_OVERRIDE { return true; } + int getLayouts(const std::vector& actualInputs, + std::vector& desiredInputs, + const int requiredOutputs, + std::vector& outputs) const CV_OVERRIDE + { + auto* netimpl_ = getNetImpl(this); + DataLayout defaultLayout = netimpl_->originalLayout; + const size_t ninputs = actualInputs.size(); + desiredInputs = actualInputs; + outputs.assign(requiredOutputs, DATA_LAYOUT_UNKNOWN); + + const bool inputIsBlock = ninputs >= 1 && actualInputs[0] == DATA_LAYOUT_BLOCK; + + std::vector resolvedAxes = axes; + if (resolvedAxes.empty() && this->inputs.size() > 3 && + netimpl_->isConstArg(this->inputs[3])) { + Mat axesT = netimpl_->argTensor(this->inputs[3]); + tensorToIntVec(axesT, resolvedAxes); + } + bool axesOK = !resolvedAxes.empty(); + if (axesOK) { + const int channelAxis = (defaultLayout == DATA_LAYOUT_NCHW) ? 1 : + (defaultLayout == DATA_LAYOUT_NHWC) ? 3 : -1; + for (int a : resolvedAxes) { + if (a < 0 || a == channelAxis) { axesOK = false; break; } + } + } + + if (inputIsBlock && axesOK) { + outputs.assign(requiredOutputs, DATA_LAYOUT_BLOCK); + } else if (inputIsBlock) { + desiredInputs[0] = defaultLayout; + } + return outputs[0] == DATA_LAYOUT_BLOCK ? netimpl_->defaultC0 : 0; + } + bool getMemoryShapes(const std::vector &inputs, const int, std::vector &outputs, diff --git a/modules/dnn/src/layers/transform_layout_layer.cpp b/modules/dnn/src/layers/transform_layout_layer.cpp index 08fba700ff..21a3b6ed3d 100644 --- a/modules/dnn/src/layers/transform_layout_layer.cpp +++ b/modules/dnn/src/layers/transform_layout_layer.cpp @@ -6,6 +6,10 @@ #include "layers_common.hpp" #include "../net_impl.hpp" +#if defined(__AVX2__) +#include +#endif + namespace cv { namespace dnn @@ -38,6 +42,51 @@ template static inline void transpose8x8(const _Tp* inp_, size_t istep, _Tp* out_, size_t ostep) { +#if defined(__AVX2__) + if (sizeof(_Tp) == 4u) { + // 8x8 32-bit transpose via 256-bit AVX2: 8 unpack + 8 shuffle + 8 perm. + // Roughly 2x faster than the four v_transpose4x4 path on AVX2 hosts. + const float* inp = (const float*)inp_; + float* out = (float*)out_; + __m256 r0 = _mm256_loadu_ps(inp + istep*0); + __m256 r1 = _mm256_loadu_ps(inp + istep*1); + __m256 r2 = _mm256_loadu_ps(inp + istep*2); + __m256 r3 = _mm256_loadu_ps(inp + istep*3); + __m256 r4 = _mm256_loadu_ps(inp + istep*4); + __m256 r5 = _mm256_loadu_ps(inp + istep*5); + __m256 r6 = _mm256_loadu_ps(inp + istep*6); + __m256 r7 = _mm256_loadu_ps(inp + istep*7); + + __m256 t0 = _mm256_unpacklo_ps(r0, r1); + __m256 t1 = _mm256_unpackhi_ps(r0, r1); + __m256 t2 = _mm256_unpacklo_ps(r2, r3); + __m256 t3 = _mm256_unpackhi_ps(r2, r3); + __m256 t4 = _mm256_unpacklo_ps(r4, r5); + __m256 t5 = _mm256_unpackhi_ps(r4, r5); + __m256 t6 = _mm256_unpacklo_ps(r6, r7); + __m256 t7 = _mm256_unpackhi_ps(r6, r7); + + __m256 v0 = _mm256_shuffle_ps(t0, t2, _MM_SHUFFLE(1,0,1,0)); + __m256 v1 = _mm256_shuffle_ps(t0, t2, _MM_SHUFFLE(3,2,3,2)); + __m256 v2 = _mm256_shuffle_ps(t1, t3, _MM_SHUFFLE(1,0,1,0)); + __m256 v3 = _mm256_shuffle_ps(t1, t3, _MM_SHUFFLE(3,2,3,2)); + __m256 v4 = _mm256_shuffle_ps(t4, t6, _MM_SHUFFLE(1,0,1,0)); + __m256 v5 = _mm256_shuffle_ps(t4, t6, _MM_SHUFFLE(3,2,3,2)); + __m256 v6 = _mm256_shuffle_ps(t5, t7, _MM_SHUFFLE(1,0,1,0)); + __m256 v7 = _mm256_shuffle_ps(t5, t7, _MM_SHUFFLE(3,2,3,2)); + + // 0x20 -> {lo of A, lo of B}; 0x31 -> {hi of A, hi of B} + _mm256_storeu_ps(out + ostep*0, _mm256_permute2f128_ps(v0, v4, 0x20)); + _mm256_storeu_ps(out + ostep*1, _mm256_permute2f128_ps(v1, v5, 0x20)); + _mm256_storeu_ps(out + ostep*2, _mm256_permute2f128_ps(v2, v6, 0x20)); + _mm256_storeu_ps(out + ostep*3, _mm256_permute2f128_ps(v3, v7, 0x20)); + _mm256_storeu_ps(out + ostep*4, _mm256_permute2f128_ps(v0, v4, 0x31)); + _mm256_storeu_ps(out + ostep*5, _mm256_permute2f128_ps(v1, v5, 0x31)); + _mm256_storeu_ps(out + ostep*6, _mm256_permute2f128_ps(v2, v6, 0x31)); + _mm256_storeu_ps(out + ostep*7, _mm256_permute2f128_ps(v3, v7, 0x31)); + return; + } +#endif #if CV_SIMD128 if (sizeof(_Tp) == 4u) { const uint32_t* inp = (const uint32_t*)inp_; @@ -257,10 +306,14 @@ void transformLayout(const Mat& inp, Mat& out, } size_t total = N*C1*planesize*C0; - constexpr size_t min_elems_per_chunk = 1 << 17; + // Tuned to keep small encoder/decoder transforms (e.g. 256ch * 14*14 ~ 50K elems) + // from running single-threaded. 16K elems ~ 64 KB ~ L1-resident chunk. + constexpr size_t min_elems_per_chunk = 1 << 14; int nblocks = int((total + min_elems_per_chunk/2) / min_elems_per_chunk); - nblocks = clamp(nblocks, 1, 128); + int nthreads = std::max(1, getNumThreads()); + nblocks = clamp(nblocks, 1, std::max(N*C1, 1) * 4); nblocks = (nblocks + N*C1 - 1)/(N*C1); + nblocks = std::min(nblocks, std::max(1, nthreads)); parallel_for_(Range(0, N*C1*nblocks), [&](const Range& range) { diff --git a/modules/dnn/src/net_impl.hpp b/modules/dnn/src/net_impl.hpp index 832be12cf1..eb60202932 100644 --- a/modules/dnn/src/net_impl.hpp +++ b/modules/dnn/src/net_impl.hpp @@ -456,6 +456,9 @@ struct Net::Impl : public detail::NetImplBase void fuseBasic(); // fuse ViT-style multi-head attention subgraphs void fuseAttention(); + // rewrite MatMul(A, const_B [, const_bias]) into Gemm so projection-style + // matmuls reach the MLAS pre-packed sgemm path + void fuseMatMulConstBToGemm(); // fuse Gemm layers that share the same input into one wider Gemm void fuseSharedInputGemm(); // collapse redundant Reshape/Transpose chains diff --git a/modules/dnn/src/net_impl2.cpp b/modules/dnn/src/net_impl2.cpp index 42bce5e8f8..773d8512e0 100644 --- a/modules/dnn/src/net_impl2.cpp +++ b/modules/dnn/src/net_impl2.cpp @@ -550,6 +550,7 @@ void Net::Impl::prepareForInference() fuseBN(); constArgs(); fuseAttention(); + fuseMatMulConstBToGemm(); fuseSharedInputGemm(); fuseReshapeTranspose(); fuseTransposeMatMul(); diff --git a/modules/dnn/test/test_graph_simplifier.cpp b/modules/dnn/test/test_graph_simplifier.cpp index 83f112dfe5..59d527acef 100644 --- a/modules/dnn/test/test_graph_simplifier.cpp +++ b/modules/dnn/test/test_graph_simplifier.cpp @@ -153,7 +153,10 @@ TEST_F(Test_Graph_Simplifier, BiasedMatMulSubgraph) { /* Test for 1 subgraphs - BiasedMatMulSubgraph */ - test("biased_matmul", "MatMul"); + auto engine_forced = static_cast( + cv::utils::getConfigurationParameterSizeT("OPENCV_FORCE_DNN_ENGINE", cv::dnn::ENGINE_AUTO)); + const std::string expected = engine_forced == cv::dnn::ENGINE_CLASSIC ? "MatMul" : "Gemm"; + test("biased_matmul", expected); } }}