mirror of
https://github.com/opencv/opencv.git
synced 2026-07-29 15:23:05 +04:00
Merge pull request #29126 from abhishek-gola:flash_attention
Attention graph fusion with MLAS FlashAttention #29126 Performance numbers for Owl-v2 model on intel i9: ``` ORT: Average inference time over 10 runs: 1411.55 ms (min 1399.75, max 1438.89) NEW: Average inference time over 10 runs: 1078 ms (min 1048.04, max 1110.61) ``` ### Pull Request Readiness Checklist See details at https://github.com/opencv/opencv/wiki/How_to_contribute#making-a-good-pull-request - [x] I agree to contribute to the project under Apache 2 License. - [x] To the best of my knowledge, the proposed patch is not based on a code under GPL or another license that is incompatible with OpenCV - [x] The PR is proposed to the proper branch - [x] There is a reference to the original bug report and related work - [x] There is accuracy test, performance test and test data in opencv_extra repository, if applicable Patch to opencv_extra has the same branch name. - [x] The feature is well documented and sample code can be built with the project CMake
This commit is contained in:
@@ -309,6 +309,499 @@ struct ModelFusionAttention
|
||||
return transpose_idx;
|
||||
}
|
||||
|
||||
// Combined-QKV attention: QKV proj -> Reshape ->
|
||||
// Transpose -> 3 Gathers -> QK^T -> Softmax(no mask) -> *V.
|
||||
bool tryFuseCombinedQKV(const vector<Ptr<Layer>>& prog, int qkv_matmul_idx,
|
||||
std::set<int>& removed_ops,
|
||||
vector<std::pair<int, Ptr<Layer>>>& replacements)
|
||||
{
|
||||
if (qkv_matmul_idx < 0 || qkv_matmul_idx >= (int)prog.size() || !prog[qkv_matmul_idx])
|
||||
return false;
|
||||
if (removed_ops.count(qkv_matmul_idx)) return false;
|
||||
if (!isProjCandidate(prog[qkv_matmul_idx])) return false;
|
||||
|
||||
Mat W = getProjWeight(prog[qkv_matmul_idx]);
|
||||
if (W.dims != 2 || W.size[1] % 3 != 0) return false;
|
||||
const int input_hidden = W.size[0];
|
||||
const int proj_hidden = W.size[1] / 3;
|
||||
const int total_hidden = W.size[1];
|
||||
|
||||
Arg cur = prog[qkv_matmul_idx]->outputs[0];
|
||||
|
||||
// Optional bias: separate Add op or Gemm's blobs[1].
|
||||
int add_idx = -1;
|
||||
Mat bias_mat;
|
||||
bool has_bias = false;
|
||||
if (prog[qkv_matmul_idx]->blobs.size() >= 2) {
|
||||
bias_mat = prog[qkv_matmul_idx]->blobs[1];
|
||||
has_bias = bias_mat.total() == (size_t)total_hidden && bias_mat.type() == CV_32F;
|
||||
if (!has_bias) return false;
|
||||
}
|
||||
int next = singleConsumer(cur);
|
||||
if (!has_bias && next >= 0 && next < (int)prog.size() && prog[next]) {
|
||||
NaryEltwiseLayer* e = dynamic_cast<NaryEltwiseLayer*>(prog[next].get());
|
||||
if (e && e->op == NaryEltwiseLayer::OPERATION::ADD &&
|
||||
prog[next]->inputs.size() == 2)
|
||||
{
|
||||
Mat b_candidate;
|
||||
for (Arg in : prog[next]->inputs) {
|
||||
if (netimpl->isConstArg(in)) {
|
||||
Mat t = netimpl->argTensor(in);
|
||||
if (t.type() == CV_32F && (int)t.total() == total_hidden)
|
||||
b_candidate = t;
|
||||
}
|
||||
}
|
||||
if (!b_candidate.empty()) {
|
||||
add_idx = next;
|
||||
bias_mat = b_candidate;
|
||||
has_bias = true;
|
||||
cur = prog[add_idx]->outputs[0];
|
||||
next = singleConsumer(cur);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (!isReshape(prog, next)) return false;
|
||||
const int reshape_idx = next;
|
||||
const auto& rinputs = prog[reshape_idx]->inputs;
|
||||
if (rinputs.size() < 2) return false;
|
||||
|
||||
std::set<int> extra_ops;
|
||||
int num_heads = -1, head_dim = -1;
|
||||
Arg shape_arg = rinputs[1];
|
||||
if (netimpl->isConstArg(shape_arg)) {
|
||||
Mat sh = netimpl->argTensor(shape_arg);
|
||||
if (sh.total() != 5) return false;
|
||||
const int64_t* sd = sh.ptr<int64_t>();
|
||||
if ((int)sd[2] != 3) return false;
|
||||
num_heads = (int)sd[3];
|
||||
head_dim = (int)sd[4];
|
||||
} else {
|
||||
auto it = producer_.find(shape_arg.idx);
|
||||
if (it == producer_.end()) return false;
|
||||
int concat_idx = it->second;
|
||||
if (concat_idx < 0 || concat_idx >= (int)prog.size() || !prog[concat_idx])
|
||||
return false;
|
||||
if (!dynamic_cast<Concat2Layer*>(prog[concat_idx].get())) return false;
|
||||
const auto& cinputs = prog[concat_idx]->inputs;
|
||||
if (cinputs.size() != 5) return false;
|
||||
if (extractConstInt(prog, cinputs[2]) != 3) return false;
|
||||
num_heads = extractConstInt(prog, cinputs[3]);
|
||||
head_dim = extractConstInt(prog, cinputs[4]);
|
||||
if (num_heads <= 0 || head_dim <= 0) return false;
|
||||
collectShapeChain(prog, concat_idx, extra_ops);
|
||||
}
|
||||
if (num_heads * head_dim != proj_hidden) return false;
|
||||
|
||||
// Transpose must bring "3" to front; only perm[0]==2 is required.
|
||||
Arg reshape_out = prog[reshape_idx]->outputs[0];
|
||||
int transpose_idx = singleConsumer(reshape_out);
|
||||
if (!isTranspose(prog, transpose_idx)) return false;
|
||||
TransposeLayer* tr = dynamic_cast<TransposeLayer*>(prog[transpose_idx].get());
|
||||
if (!tr || tr->perm.size() != 5 || tr->perm[0] != 2) return false;
|
||||
|
||||
Arg tr_out = prog[transpose_idx]->outputs[0];
|
||||
auto cit = consumers_.find(tr_out.idx);
|
||||
if (cit == consumers_.end() || cit->second.size() != 3) return false;
|
||||
|
||||
int gather_for_idx[3] = { -1, -1, -1 }; // [Q,K,V] = [0,1,2]
|
||||
for (int c : cit->second) {
|
||||
if (c < 0 || c >= (int)prog.size() || !prog[c]) return false;
|
||||
Gather2Layer* g = dynamic_cast<Gather2Layer*>(prog[c].get());
|
||||
if (!g || g->axis != 0 || prog[c]->inputs.size() < 2) return false;
|
||||
Arg idx_arg = prog[c]->inputs[1];
|
||||
if (!netimpl->isConstArg(idx_arg)) return false;
|
||||
Mat t = netimpl->argTensor(idx_arg);
|
||||
if (t.total() != 1) return false;
|
||||
int v;
|
||||
if (t.type() == CV_64S) v = (int)t.at<int64_t>(0);
|
||||
else if (t.type() == CV_32S) v = (int)t.at<int32_t>(0);
|
||||
else return false;
|
||||
if (v < 0 || v > 2 || gather_for_idx[v] != -1) return false;
|
||||
gather_for_idx[v] = c;
|
||||
}
|
||||
const int q_gather = gather_for_idx[0];
|
||||
const int k_gather = gather_for_idx[1];
|
||||
const int v_gather = gather_for_idx[2];
|
||||
if (q_gather < 0 || k_gather < 0 || v_gather < 0) return false;
|
||||
|
||||
Arg q_out = prog[q_gather]->outputs[0];
|
||||
Arg k_out = prog[k_gather]->outputs[0];
|
||||
Arg v_out = prog[v_gather]->outputs[0];
|
||||
|
||||
int q_mul_idx = singleConsumer(q_out);
|
||||
float q_scale = 1.f;
|
||||
if (!isScalarMul(prog, q_mul_idx, &q_scale)) return false;
|
||||
if (q_scale == 0.f) return false;
|
||||
int qk_matmul_idx = singleConsumer(prog[q_mul_idx]->outputs[0]);
|
||||
if (!isMatMul(prog, qk_matmul_idx)) return false;
|
||||
|
||||
int k_trans_idx = singleConsumer(k_out);
|
||||
if (!isTranspose(prog, k_trans_idx)) return false;
|
||||
Arg k_trans_out = prog[k_trans_idx]->outputs[0];
|
||||
bool k_connected = false;
|
||||
for (Arg in : prog[qk_matmul_idx]->inputs)
|
||||
if (in.idx == k_trans_out.idx) { k_connected = true; break; }
|
||||
if (!k_connected) return false;
|
||||
|
||||
int softmax_idx = singleConsumer(prog[qk_matmul_idx]->outputs[0]);
|
||||
if (!isSoftmax(prog, softmax_idx)) return false;
|
||||
int av_matmul_idx = singleConsumer(prog[softmax_idx]->outputs[0]);
|
||||
if (!isMatMul(prog, av_matmul_idx)) return false;
|
||||
bool v_connected = false;
|
||||
for (Arg in : prog[av_matmul_idx]->inputs)
|
||||
if (in.idx == v_out.idx) { v_connected = true; break; }
|
||||
if (!v_connected) return false;
|
||||
|
||||
int out_trans_idx = singleConsumer(prog[av_matmul_idx]->outputs[0]);
|
||||
if (!isTranspose(prog, out_trans_idx)) return false;
|
||||
int out_reshape_idx = findMatchingConsumer(prog, prog[out_trans_idx]->outputs[0],
|
||||
[](Layer* L){ return dynamic_cast<Reshape2Layer*>(L) != nullptr; },
|
||||
&extra_ops);
|
||||
if (!isReshape(prog, out_reshape_idx)) return false;
|
||||
|
||||
const auto& or_inputs = prog[out_reshape_idx]->inputs;
|
||||
if (or_inputs.size() >= 2 && !netimpl->isConstArg(or_inputs[1])) {
|
||||
auto it = producer_.find(or_inputs[1].idx);
|
||||
if (it != producer_.end())
|
||||
collectShapeChain(prog, it->second, extra_ops);
|
||||
}
|
||||
|
||||
// W is already laid out as [Q|K|V] along the output dim.
|
||||
Mat W_qkv = W.clone();
|
||||
Mat bias_qkv;
|
||||
if (has_bias) bias_qkv = bias_mat.clone();
|
||||
|
||||
// Attention `scale` is the pre-softmax divisor; Q was multiplied by
|
||||
// q_scale, so scale = 1/q_scale.
|
||||
const float param_scale = 1.0f / q_scale;
|
||||
|
||||
LayerParams attn_params;
|
||||
attn_params.name = prog[qkv_matmul_idx]->name + "_fused_attention";
|
||||
attn_params.type = "Attention";
|
||||
attn_params.set("num_heads", num_heads);
|
||||
int qkv_sizes[3] = { proj_hidden, proj_hidden, proj_hidden };
|
||||
attn_params.set("qkv_hidden_sizes", DictValue::arrayInt(qkv_sizes, 3));
|
||||
attn_params.set("scale", param_scale);
|
||||
attn_params.set("output_ndims", 3);
|
||||
attn_params.blobs.push_back(W_qkv);
|
||||
if (has_bias) attn_params.blobs.push_back(bias_qkv);
|
||||
|
||||
Ptr<Layer> attn_layer = LayerFactory::createLayerInstance(attn_params.type, attn_params);
|
||||
CV_Assert(attn_layer);
|
||||
Arg shared_input = prog[qkv_matmul_idx]->inputs[0];
|
||||
attn_layer->inputs = { shared_input };
|
||||
attn_layer->outputs = prog[out_reshape_idx]->outputs;
|
||||
attn_layer->netimpl = netimpl;
|
||||
|
||||
std::set<int> to_remove = {
|
||||
qkv_matmul_idx, reshape_idx, transpose_idx,
|
||||
q_gather, k_gather, v_gather,
|
||||
q_mul_idx, k_trans_idx,
|
||||
qk_matmul_idx, softmax_idx,
|
||||
av_matmul_idx, out_trans_idx, out_reshape_idx
|
||||
};
|
||||
if (add_idx >= 0) to_remove.insert(add_idx);
|
||||
for (int op : extra_ops) to_remove.insert(op);
|
||||
|
||||
(void)input_hidden; // input_hidden currently unused; kept for symmetry with existing path.
|
||||
|
||||
for (int op : to_remove) removed_ops.insert(op);
|
||||
const int insert_pos = *std::min_element(to_remove.begin(), to_remove.end());
|
||||
replacements.push_back({insert_pos, attn_layer});
|
||||
return true;
|
||||
}
|
||||
|
||||
// CLIP-branch trace: arg -> [Transpose3D(K^T)] -> Reshape3D -> Transpose ->
|
||||
// Reshape4D -> [Mul(Q scale)] -> [Add(bias)] -> proj_MatMul.
|
||||
int traceClipBranch(const vector<Ptr<Layer>>& prog, Arg arg,
|
||||
bool is_q_branch, bool is_k_branch,
|
||||
Mat& out_W, Mat& out_bias, int& out_num_heads,
|
||||
float& out_q_scale,
|
||||
std::set<int>& ops_consumed) const
|
||||
{
|
||||
Arg cur = arg;
|
||||
|
||||
auto stepProducer = [&](Arg a) -> int {
|
||||
auto it = producer_.find(a.idx);
|
||||
return it == producer_.end() ? -1 : it->second;
|
||||
};
|
||||
|
||||
// K side: peel the (B*H,S,D) -> (B*H,D,S) transpose-3D first.
|
||||
if (is_k_branch) {
|
||||
int idx = stepProducer(cur);
|
||||
if (idx < 0 || !prog[idx]) return -1;
|
||||
TransposeLayer* tr = dynamic_cast<TransposeLayer*>(prog[idx].get());
|
||||
if (!tr || tr->perm.size() != 3) return -1;
|
||||
if (tr->perm[0] != 0 || tr->perm[1] != 2 || tr->perm[2] != 1) return -1;
|
||||
ops_consumed.insert(idx);
|
||||
cur = prog[idx]->inputs[0];
|
||||
}
|
||||
|
||||
// (B,H,S,D) -> (B*H,S,D); shape may be dynamic, validated downstream.
|
||||
int r3d_idx = stepProducer(cur);
|
||||
if (!isReshape(prog, r3d_idx)) return -1;
|
||||
ops_consumed.insert(r3d_idx);
|
||||
if (prog[r3d_idx]->inputs.size() < 2) return -1;
|
||||
Arg shape_arg_r3d = prog[r3d_idx]->inputs[1];
|
||||
if (!netimpl->isConstArg(shape_arg_r3d)) {
|
||||
int sh_idx = stepProducer(shape_arg_r3d);
|
||||
if (sh_idx >= 0)
|
||||
collectShapeChain(prog, sh_idx, ops_consumed);
|
||||
}
|
||||
cur = prog[r3d_idx]->inputs[0];
|
||||
|
||||
int t4d_idx = stepProducer(cur);
|
||||
if (!isTranspose(prog, t4d_idx)) return -1;
|
||||
TransposeLayer* tr4 = dynamic_cast<TransposeLayer*>(prog[t4d_idx].get());
|
||||
if (!tr4 || tr4->perm.size() != 4) return -1;
|
||||
if (tr4->perm[0] != 0 || tr4->perm[1] != 2 ||
|
||||
tr4->perm[2] != 1 || tr4->perm[3] != 3) return -1;
|
||||
ops_consumed.insert(t4d_idx);
|
||||
cur = prog[t4d_idx]->inputs[0];
|
||||
|
||||
// (B,S,H*D) -> (B,S,H,D); num_heads is the third shape entry.
|
||||
int r4d_idx = stepProducer(cur);
|
||||
if (!isReshape(prog, r4d_idx)) return -1;
|
||||
if (prog[r4d_idx]->inputs.size() < 2) return -1;
|
||||
Arg shape_arg_r4d = prog[r4d_idx]->inputs[1];
|
||||
int num_heads = -1;
|
||||
if (netimpl->isConstArg(shape_arg_r4d)) {
|
||||
Mat shape_mat = netimpl->argTensor(shape_arg_r4d);
|
||||
if (shape_mat.total() != 4) return -1;
|
||||
if (shape_mat.type() == CV_64S)
|
||||
num_heads = static_cast<int>(shape_mat.ptr<int64_t>()[2]);
|
||||
else if (shape_mat.type() == CV_32S)
|
||||
num_heads = static_cast<int>(shape_mat.ptr<int32_t>()[2]);
|
||||
else return -1;
|
||||
} else {
|
||||
int concat_idx = stepProducer(shape_arg_r4d);
|
||||
if (concat_idx < 0 || !prog[concat_idx]) return -1;
|
||||
if (!dynamic_cast<Concat2Layer*>(prog[concat_idx].get())) return -1;
|
||||
const auto& cinputs = prog[concat_idx]->inputs;
|
||||
if (cinputs.size() != 4) return -1;
|
||||
num_heads = extractConstInt(prog, cinputs[2]);
|
||||
if (num_heads <= 0) return -1;
|
||||
collectShapeChain(prog, concat_idx, ops_consumed);
|
||||
}
|
||||
if (num_heads <= 0) return -1;
|
||||
out_num_heads = num_heads;
|
||||
ops_consumed.insert(r4d_idx);
|
||||
cur = prog[r4d_idx]->inputs[0];
|
||||
|
||||
out_q_scale = 1.0f;
|
||||
if (is_q_branch) {
|
||||
int mul_idx = stepProducer(cur);
|
||||
if (mul_idx < 0 || !prog[mul_idx]) return -1;
|
||||
NaryEltwiseLayer* mul =
|
||||
dynamic_cast<NaryEltwiseLayer*>(prog[mul_idx].get());
|
||||
if (!mul || mul->op != NaryEltwiseLayer::OPERATION::PROD) return -1;
|
||||
if (prog[mul_idx]->inputs.size() != 2) return -1;
|
||||
Arg runtime_arg;
|
||||
bool got_scale = false, got_runtime = false;
|
||||
for (Arg in : prog[mul_idx]->inputs) {
|
||||
if (netimpl->isConstArg(in)) {
|
||||
Mat t = netimpl->argTensor(in);
|
||||
if (t.total() != 1) return -1;
|
||||
if (t.type() == CV_32F) out_q_scale = t.at<float>(0);
|
||||
else if (t.type() == CV_64F) out_q_scale = (float)t.at<double>(0);
|
||||
else return -1;
|
||||
got_scale = true;
|
||||
} else {
|
||||
runtime_arg = in;
|
||||
got_runtime = true;
|
||||
}
|
||||
}
|
||||
if (!got_scale || !got_runtime) return -1;
|
||||
ops_consumed.insert(mul_idx);
|
||||
cur = runtime_arg;
|
||||
}
|
||||
|
||||
int next_idx = stepProducer(cur);
|
||||
if (next_idx < 0 || !prog[next_idx]) return -1;
|
||||
|
||||
out_bias = Mat();
|
||||
int mm_idx = -1;
|
||||
if (dynamic_cast<NaryEltwiseLayer*>(prog[next_idx].get())) {
|
||||
NaryEltwiseLayer* add =
|
||||
dynamic_cast<NaryEltwiseLayer*>(prog[next_idx].get());
|
||||
if (!add || add->op != NaryEltwiseLayer::OPERATION::ADD) return -1;
|
||||
if (prog[next_idx]->inputs.size() != 2) return -1;
|
||||
Arg bias_arg;
|
||||
Arg matmul_out_arg;
|
||||
bool got_bias = false, got_runtime2 = false;
|
||||
for (Arg in : prog[next_idx]->inputs) {
|
||||
if (netimpl->isConstArg(in)) {
|
||||
bias_arg = in;
|
||||
got_bias = true;
|
||||
} else {
|
||||
matmul_out_arg = in;
|
||||
got_runtime2 = true;
|
||||
}
|
||||
}
|
||||
if (!got_bias || !got_runtime2) return -1;
|
||||
out_bias = netimpl->argTensor(bias_arg).clone();
|
||||
ops_consumed.insert(next_idx);
|
||||
mm_idx = stepProducer(matmul_out_arg);
|
||||
} else {
|
||||
mm_idx = next_idx;
|
||||
}
|
||||
|
||||
if (mm_idx < 0 || !prog[mm_idx]) return -1;
|
||||
if (!dynamic_cast<MatMulLayer*>(prog[mm_idx].get())) return -1;
|
||||
if (prog[mm_idx]->blobs.empty()) return -1;
|
||||
if (prog[mm_idx]->inputs.size() != 1) return -1;
|
||||
out_W = prog[mm_idx]->blobs[0].clone();
|
||||
// Folded MatMul carries bias as a second blob (real_ndims_C >= 1).
|
||||
if (out_bias.empty() && prog[mm_idx]->blobs.size() >= 2)
|
||||
out_bias = prog[mm_idx]->blobs.back().clone();
|
||||
return mm_idx;
|
||||
}
|
||||
|
||||
// 3 separate q/k/v projections, Q scaled, K^T at
|
||||
// the QK^T matmul, output reshaped+transposed back to (B,S,H*D).
|
||||
bool tryFuseClipAttention(const vector<Ptr<Layer>>& prog, int softmax_idx,
|
||||
std::set<int>& removed_ops,
|
||||
vector<std::pair<int, Ptr<Layer>>>& replacements)
|
||||
{
|
||||
if (softmax_idx < 0 || softmax_idx >= (int)prog.size() || !prog[softmax_idx])
|
||||
return false;
|
||||
if (removed_ops.count(softmax_idx)) return false;
|
||||
SoftmaxLayer* sm = dynamic_cast<SoftmaxLayer*>(prog[softmax_idx].get());
|
||||
if (!sm || sm->logSoftMax) return false;
|
||||
if (prog[softmax_idx]->inputs.size() != 1) return false;
|
||||
|
||||
Arg sm_in = prog[softmax_idx]->inputs[0];
|
||||
auto it = producer_.find(sm_in.idx);
|
||||
if (it == producer_.end()) return false;
|
||||
int qk_matmul_idx = it->second;
|
||||
if (qk_matmul_idx < 0 || !prog[qk_matmul_idx]) return false;
|
||||
if (!dynamic_cast<MatMulLayer*>(prog[qk_matmul_idx].get())) return false;
|
||||
if (!prog[qk_matmul_idx]->blobs.empty()) return false;
|
||||
if (prog[qk_matmul_idx]->inputs.size() != 2) return false;
|
||||
Arg q_arg = prog[qk_matmul_idx]->inputs[0];
|
||||
Arg k_arg = prog[qk_matmul_idx]->inputs[1];
|
||||
|
||||
Arg sm_out = prog[softmax_idx]->outputs[0];
|
||||
int av_matmul_idx = singleConsumer(sm_out);
|
||||
if (av_matmul_idx < 0 || !prog[av_matmul_idx]) return false;
|
||||
if (!dynamic_cast<MatMulLayer*>(prog[av_matmul_idx].get())) return false;
|
||||
if (!prog[av_matmul_idx]->blobs.empty()) return false;
|
||||
if (prog[av_matmul_idx]->inputs.size() != 2) return false;
|
||||
Arg v_arg;
|
||||
bool got_v = false;
|
||||
for (Arg in : prog[av_matmul_idx]->inputs) {
|
||||
if (in.idx == sm_out.idx) continue;
|
||||
v_arg = in; got_v = true;
|
||||
}
|
||||
if (!got_v) return false;
|
||||
|
||||
std::set<int> consumed;
|
||||
Mat Wq, Wk, Wv, bq, bk, bv;
|
||||
int nh_q = 0, nh_k = 0, nh_v = 0;
|
||||
float q_scale = 1.f;
|
||||
float dummy = 1.f;
|
||||
int q_mm_idx = traceClipBranch(prog, q_arg, /*is_q=*/true, /*is_k=*/false,
|
||||
Wq, bq, nh_q, q_scale, consumed);
|
||||
if (q_mm_idx < 0) return false;
|
||||
int k_mm_idx = traceClipBranch(prog, k_arg, /*is_q=*/false, /*is_k=*/true,
|
||||
Wk, bk, nh_k, dummy, consumed);
|
||||
if (k_mm_idx < 0) return false;
|
||||
int v_mm_idx = traceClipBranch(prog, v_arg, /*is_q=*/false, /*is_k=*/false,
|
||||
Wv, bv, nh_v, dummy, consumed);
|
||||
if (v_mm_idx < 0) return false;
|
||||
|
||||
if (q_mm_idx == k_mm_idx || k_mm_idx == v_mm_idx || q_mm_idx == v_mm_idx)
|
||||
return false;
|
||||
if (nh_q != nh_k || nh_k != nh_v) return false;
|
||||
const int num_heads = nh_q;
|
||||
if (q_scale == 0.f) return false;
|
||||
|
||||
// All three projections share the same input.
|
||||
Arg shared_input = prog[q_mm_idx]->inputs[0];
|
||||
if (prog[k_mm_idx]->inputs[0].idx != shared_input.idx) return false;
|
||||
if (prog[v_mm_idx]->inputs[0].idx != shared_input.idx) return false;
|
||||
|
||||
if (Wq.dims != 2 || Wk.dims != 2 || Wv.dims != 2) return false;
|
||||
if (Wq.size[0] != Wk.size[0] || Wk.size[0] != Wv.size[0]) return false;
|
||||
const int hidden_in = Wq.size[0];
|
||||
const int q_hidden = Wq.size[1];
|
||||
const int k_hidden = Wk.size[1];
|
||||
const int v_hidden = Wv.size[1];
|
||||
if (q_hidden != k_hidden || k_hidden != v_hidden) return false;
|
||||
if (q_hidden % num_heads != 0) return false;
|
||||
|
||||
// Output chain: Reshape4D -> Transpose([0,2,1,3]) -> Reshape3D, the
|
||||
// boundary of the fused block.
|
||||
Arg av_out = prog[av_matmul_idx]->outputs[0];
|
||||
int out_r4d = singleConsumer(av_out);
|
||||
if (!isReshape(prog, out_r4d)) return false;
|
||||
Arg out_r4d_out = prog[out_r4d]->outputs[0];
|
||||
int out_t4d = singleConsumer(out_r4d_out);
|
||||
if (!isTranspose(prog, out_t4d)) return false;
|
||||
TransposeLayer* out_tr =
|
||||
dynamic_cast<TransposeLayer*>(prog[out_t4d].get());
|
||||
if (!out_tr || out_tr->perm.size() != 4) return false;
|
||||
if (out_tr->perm[0] != 0 || out_tr->perm[1] != 2 ||
|
||||
out_tr->perm[2] != 1 || out_tr->perm[3] != 3) return false;
|
||||
Arg out_t4d_out = prog[out_t4d]->outputs[0];
|
||||
int out_r3d = singleConsumer(out_t4d_out);
|
||||
if (!isReshape(prog, out_r3d)) return false;
|
||||
|
||||
// Combined [Q|K|V] weight along the output dim.
|
||||
int total_hidden = q_hidden + k_hidden + v_hidden;
|
||||
int wshape[] = {hidden_in, total_hidden};
|
||||
Mat W_qkv(2, wshape, CV_32F);
|
||||
for (int r = 0; r < hidden_in; r++) {
|
||||
float* dst = W_qkv.ptr<float>(r);
|
||||
std::memcpy(dst, Wq.ptr<float>(r), q_hidden * sizeof(float));
|
||||
std::memcpy(dst + q_hidden, Wk.ptr<float>(r), k_hidden * sizeof(float));
|
||||
std::memcpy(dst + q_hidden + k_hidden, Wv.ptr<float>(r), v_hidden * sizeof(float));
|
||||
}
|
||||
Mat bias_qkv;
|
||||
if (!bq.empty() && !bk.empty() && !bv.empty()) {
|
||||
const int bias_total = q_hidden + k_hidden + v_hidden;
|
||||
bias_qkv.create(1, &bias_total, CV_32F);
|
||||
float* dst = bias_qkv.ptr<float>();
|
||||
std::memcpy(dst, bq.ptr<float>(), q_hidden * sizeof(float));
|
||||
std::memcpy(dst + q_hidden, bk.ptr<float>(), k_hidden * sizeof(float));
|
||||
std::memcpy(dst + q_hidden + k_hidden, bv.ptr<float>(), v_hidden * sizeof(float));
|
||||
}
|
||||
|
||||
// scale (pre-softmax divisor) = 1/q_scale, since Q was pre-multiplied.
|
||||
const float param_scale = 1.f / q_scale;
|
||||
|
||||
LayerParams attn_params;
|
||||
attn_params.name = prog[q_mm_idx]->name + "_fused_attention";
|
||||
attn_params.type = "Attention";
|
||||
attn_params.set("num_heads", num_heads);
|
||||
int qkv_sizes[3] = { q_hidden, k_hidden, v_hidden };
|
||||
attn_params.set("qkv_hidden_sizes", DictValue::arrayInt(qkv_sizes, 3));
|
||||
attn_params.set("scale", param_scale);
|
||||
attn_params.set("output_ndims", 3);
|
||||
attn_params.blobs.push_back(W_qkv);
|
||||
if (!bias_qkv.empty()) attn_params.blobs.push_back(bias_qkv);
|
||||
|
||||
Ptr<Layer> attn_layer =
|
||||
LayerFactory::createLayerInstance(attn_params.type, attn_params);
|
||||
if (!attn_layer) return false;
|
||||
attn_layer->inputs = { shared_input };
|
||||
attn_layer->outputs = prog[out_r3d]->outputs;
|
||||
attn_layer->netimpl = netimpl;
|
||||
|
||||
std::set<int> to_remove = {
|
||||
q_mm_idx, k_mm_idx, v_mm_idx,
|
||||
qk_matmul_idx, softmax_idx, av_matmul_idx,
|
||||
out_r4d, out_t4d, out_r3d
|
||||
};
|
||||
for (int op : consumed) to_remove.insert(op);
|
||||
|
||||
for (int op : to_remove) removed_ops.insert(op);
|
||||
int insert_pos = *std::min_element(to_remove.begin(), to_remove.end());
|
||||
replacements.push_back({insert_pos, attn_layer});
|
||||
return true;
|
||||
}
|
||||
|
||||
bool fuseGraph(Ptr<Graph>& graph)
|
||||
{
|
||||
const vector<Ptr<Layer>>& prog = graph->prog();
|
||||
@@ -324,17 +817,25 @@ struct ModelFusionAttention
|
||||
consumers_[inp.idx].push_back((int)i);
|
||||
}
|
||||
|
||||
bool modified = false;
|
||||
std::set<int> removed_ops;
|
||||
|
||||
// Pass 1: combined-QKV blocks.
|
||||
for (size_t i = 0; i < nops; i++) {
|
||||
if (!prog[i] || removed_ops.count((int)i)) continue;
|
||||
if (tryFuseCombinedQKV(prog, (int)i, removed_ops, attention_replacements_))
|
||||
modified = true;
|
||||
}
|
||||
|
||||
// Pass 2: the original 3-separate-projection path.
|
||||
std::map<int, vector<int>> qkv_candidates;
|
||||
for (size_t i = 0; i < nops; i++) {
|
||||
if (!prog[i]) continue;
|
||||
if (!prog[i] || removed_ops.count((int)i)) continue;
|
||||
if (!isProjCandidate(prog[i])) continue;
|
||||
Arg inp = prog[i]->inputs[0];
|
||||
qkv_candidates[inp.idx].push_back((int)i);
|
||||
}
|
||||
|
||||
bool modified = false;
|
||||
std::set<int> removed_ops;
|
||||
|
||||
for (auto& candidate : qkv_candidates) {
|
||||
auto matmul_indices = candidate.second;
|
||||
if (matmul_indices.size() < 3) continue;
|
||||
@@ -372,8 +873,6 @@ struct ModelFusionAttention
|
||||
if (num_heads[0] != num_heads[1] || num_heads[1] != num_heads[2])
|
||||
continue;
|
||||
|
||||
// K is identified by perm [0,2,3,1] (transposes head_dim to second-last);
|
||||
// Q and V use perm [0,2,1,3].
|
||||
int k_slot = -1;
|
||||
vector<int> perm_k = {0, 2, 3, 1};
|
||||
for (int k = 0; k < 3; k++)
|
||||
@@ -648,6 +1147,15 @@ struct ModelFusionAttention
|
||||
}
|
||||
}
|
||||
|
||||
// Pass 3: anchored on each Softmax.
|
||||
for (size_t i = 0; i < nops; i++) {
|
||||
if (!prog[i] || removed_ops.count((int)i)) continue;
|
||||
if (prog[i]->type != "Softmax") continue;
|
||||
if (tryFuseClipAttention(prog, (int)i, removed_ops,
|
||||
attention_replacements_))
|
||||
modified = true;
|
||||
}
|
||||
|
||||
if (modified) {
|
||||
vector<Ptr<Layer>> newprog;
|
||||
std::sort(attention_replacements_.begin(), attention_replacements_.end(),
|
||||
|
||||
@@ -568,5 +568,4 @@ Ptr<AttentionLayer> AttentionLayer::create(const LayerParams ¶ms) {
|
||||
return makePtr<AttentionLayerImpl>(params);
|
||||
}
|
||||
|
||||
|
||||
}} // cv::dnn
|
||||
|
||||
@@ -128,6 +128,16 @@ public:
|
||||
}
|
||||
}
|
||||
|
||||
// BLOCK axis=1: per-input c-block sum overcounts on misaligned inputs
|
||||
// (YOLOX 80+4+1: 12 vs correct ceil(85/8)=11).
|
||||
if (outShape.layout == DATA_LAYOUT_BLOCK && axis_ == 1) {
|
||||
int total_C = 0;
|
||||
for (size_t i = 0; i < ninputs; i++) total_C += inpShapes[i].C;
|
||||
int C0 = outShape[outShape.dims - 1];
|
||||
outShape[1] = (total_C + C0 - 1) / C0;
|
||||
outShape.C = total_C;
|
||||
}
|
||||
|
||||
return outShape;
|
||||
}
|
||||
|
||||
@@ -173,7 +183,7 @@ public:
|
||||
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;
|
||||
const bool canKeepBlock = allBlock && axis >= 0;
|
||||
|
||||
if (canKeepBlock) {
|
||||
outputs.assign(requiredOutputs, DATA_LAYOUT_BLOCK);
|
||||
@@ -237,8 +247,50 @@ public:
|
||||
|
||||
void runOp(const std::vector<Mat>& inps, Mat& out, int axis_)
|
||||
{
|
||||
// Byte-level BLOCK concat needs every input's C to be a multiple of
|
||||
// C0; misaligned inputs (YOLOX 80+4+1) route through NCHW.
|
||||
if (axis_ == 1 && out.size.layout == DATA_LAYOUT_BLOCK) {
|
||||
const int C0 = out.size[out.dims - 1];
|
||||
bool allAligned = true;
|
||||
for (const auto& inp : inps) {
|
||||
if (inp.size.layout != DATA_LAYOUT_BLOCK || (inp.size.channels() % C0) != 0) {
|
||||
allAligned = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (!allAligned) {
|
||||
runOpBlockAxis1Misaligned(inps, out);
|
||||
return;
|
||||
}
|
||||
}
|
||||
concat(inps, out, axis_);
|
||||
}
|
||||
|
||||
// Fallback for axis=1 BLOCK concat when inputs aren't C0-aligned.
|
||||
void runOpBlockAxis1Misaligned(const std::vector<Mat>& inps, Mat& out)
|
||||
{
|
||||
auto* netimpl_ = getNetImpl(this);
|
||||
DataLayout origLayout = netimpl_->originalLayout;
|
||||
const int C0 = out.size[out.dims - 1];
|
||||
|
||||
std::vector<Mat> nchwInps(inps.size());
|
||||
for (size_t i = 0; i < inps.size(); i++) {
|
||||
transformLayout(inps[i], nchwInps[i], origLayout, origLayout, C0);
|
||||
}
|
||||
|
||||
// NCHW output shape: drop the trailing C0 dim, use logical C from out.size.C.
|
||||
const int dims = out.dims;
|
||||
std::vector<int> nchwDims(dims - 1);
|
||||
nchwDims[0] = out.size[0];
|
||||
nchwDims[1] = out.size.C;
|
||||
for (int d = 2; d < dims - 1; d++) nchwDims[d] = out.size[d];
|
||||
MatShape nchwShape(nchwDims, origLayout, out.size.C);
|
||||
Mat nchwOut;
|
||||
nchwOut.fit(nchwShape, out.type());
|
||||
|
||||
concat(nchwInps, nchwOut, /*axis=*/1);
|
||||
transformLayout(nchwOut, out, DATA_LAYOUT_BLOCK, origLayout, C0);
|
||||
}
|
||||
};
|
||||
|
||||
Ptr<Concat2Layer> Concat2Layer::create(const LayerParams& params)
|
||||
|
||||
@@ -6,7 +6,10 @@
|
||||
#include "../net_impl.hpp"
|
||||
#include "layers_common.hpp"
|
||||
#include "conv2_common.hpp"
|
||||
#include "cpu_kernels/mlas_gemm.hpp"
|
||||
#include "opencv2/core/hal/intrin.hpp"
|
||||
#include <algorithm>
|
||||
#include <cstring>
|
||||
|
||||
namespace cv
|
||||
{
|
||||
@@ -124,6 +127,41 @@ public:
|
||||
CV_Assert(bias_.isContinuous() && bias_.total() == wshape0[0]);
|
||||
bias_.convertTo(bias, CV_32F);
|
||||
}
|
||||
|
||||
// Pre-pack for MLAS 1x1 SGEMM (1x1 dense, stride-1, FP32, Cout*Cin
|
||||
// >= 256*256 so the reorder cost is amortized).
|
||||
mlas_packed_B_.release();
|
||||
mlas_packed_M_ = mlas_packed_K_ = 0;
|
||||
if (!depthwise && ngroups == 1 && wtype0 == CV_32F &&
|
||||
wtype == CV_32F && mlasAvailable())
|
||||
{
|
||||
bool ksize_all_one = wshape0.dims >= 3;
|
||||
for (int k = 2; k < wshape0.dims; k++)
|
||||
if (wshape0[k] != 1) { ksize_all_one = false; break; }
|
||||
bool strides_all_one = !strides.empty();
|
||||
for (int s : strides) if (s != 1) { strides_all_one = false; break; }
|
||||
const int Cout = wshape0[0];
|
||||
const int Cin = wshape0[1];
|
||||
const bool big_enough = (int64_t)Cout * (int64_t)Cin >= (int64_t)(256 * 256);
|
||||
if (ksize_all_one && strides_all_one && big_enough) {
|
||||
size_t pack_bytes = mlasSgemmPackBSize(false, true, Cout, Cin);
|
||||
if (pack_bytes > 0) {
|
||||
mlas_packed_B_.create(1, (int)pack_bytes, CV_8U);
|
||||
// Weight is (Cout, Cin, 1, ..., 1) contiguous; reshape to
|
||||
// (Cout, Cin) and PackB with trans_b=true.
|
||||
Mat W2D = weights_.reshape(1, Cout);
|
||||
if (!mlasSgemmPackB(false, true, Cout, Cin,
|
||||
W2D.ptr<float>(), Cin,
|
||||
mlas_packed_B_.data))
|
||||
{
|
||||
mlas_packed_B_.release();
|
||||
} else {
|
||||
mlas_packed_M_ = Cout;
|
||||
mlas_packed_K_ = Cin;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void fuseBatchNormWeights(const BatchNorm2Layer* bn)
|
||||
@@ -386,6 +424,23 @@ public:
|
||||
void* outptr = out.data;
|
||||
const void* wptr = weights.data;
|
||||
|
||||
// MLAS 1x1 SGEMM path. Skipped for small spatial: the gather/scatter
|
||||
// tax exceeds MLAS's SGEMM speedup over the in-place NCHWc8 kernel.
|
||||
if (mlas1x1Enabled() && inptype == CV_32F &&
|
||||
!activationFunc && !addResidual && inpshape.back() == 8)
|
||||
{
|
||||
const int ndims = inpshape.dims;
|
||||
int HW = 1;
|
||||
for (int i = 2; i < ndims - 1; i++) HW *= inpshape[i];
|
||||
constexpr int MLAS_MIN_SPATIAL = 256;
|
||||
if (HW >= MLAS_MIN_SPATIAL) {
|
||||
forwardMlas1x1(inp, out);
|
||||
if (uouts) out.copyTo(uouts->at(0));
|
||||
if (dynamicWeights) weights.release();
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
ConvFunc func = cs.depthwise ? getDepthwiseConvFunc(inptype) : getConvFunc(inptype, C0);
|
||||
CV_Assert(func != nullptr);
|
||||
func(inptr, resptr, outptr, cs, wptr, scale_data, bias_data);
|
||||
@@ -402,6 +457,206 @@ public:
|
||||
}
|
||||
}
|
||||
|
||||
// True iff setWeights() armed the MLAS 1x1 SGEMM path for this conv.
|
||||
bool mlas1x1Enabled() const {
|
||||
return !mlas_packed_B_.empty() && !addResidual && mlasAvailable();
|
||||
}
|
||||
|
||||
// NCHWc8 -> row-major (M_chunk, Cin); (c1-outer, mi-inner) so each
|
||||
// thread streams one C-block.
|
||||
static void gatherNCHWcToNHWC(const float* inp, int C1, int HW,
|
||||
int p_start, int p_end,
|
||||
float* dst, int Cin)
|
||||
{
|
||||
const int M = p_end - p_start;
|
||||
cv::parallel_for_(cv::Range(0, C1), [&](const cv::Range& r) {
|
||||
for (int c1 = r.start; c1 < r.end; c1++) {
|
||||
const int c_base = c1 * 8;
|
||||
const int n_valid_in = std::min(8, Cin - c_base);
|
||||
const float* src = inp + (size_t)c1 * (size_t)HW * 8
|
||||
+ (size_t)p_start * 8;
|
||||
float* dst_col = dst + c_base;
|
||||
int mi = 0;
|
||||
if (n_valid_in == 8) {
|
||||
#if CV_SIMD256 && defined(__AVX2__)
|
||||
for (; mi + 8 <= M; mi += 8) {
|
||||
__m256 r0 = _mm256_loadu_ps(src + 0 * 8);
|
||||
__m256 r1 = _mm256_loadu_ps(src + 1 * 8);
|
||||
__m256 r2 = _mm256_loadu_ps(src + 2 * 8);
|
||||
__m256 r3 = _mm256_loadu_ps(src + 3 * 8);
|
||||
__m256 r4 = _mm256_loadu_ps(src + 4 * 8);
|
||||
__m256 r5 = _mm256_loadu_ps(src + 5 * 8);
|
||||
__m256 r6 = _mm256_loadu_ps(src + 6 * 8);
|
||||
__m256 r7 = _mm256_loadu_ps(src + 7 * 8);
|
||||
_mm256_storeu_ps(dst_col + (size_t)(mi + 0) * Cin, r0);
|
||||
_mm256_storeu_ps(dst_col + (size_t)(mi + 1) * Cin, r1);
|
||||
_mm256_storeu_ps(dst_col + (size_t)(mi + 2) * Cin, r2);
|
||||
_mm256_storeu_ps(dst_col + (size_t)(mi + 3) * Cin, r3);
|
||||
_mm256_storeu_ps(dst_col + (size_t)(mi + 4) * Cin, r4);
|
||||
_mm256_storeu_ps(dst_col + (size_t)(mi + 5) * Cin, r5);
|
||||
_mm256_storeu_ps(dst_col + (size_t)(mi + 6) * Cin, r6);
|
||||
_mm256_storeu_ps(dst_col + (size_t)(mi + 7) * Cin, r7);
|
||||
src += 8 * 8;
|
||||
}
|
||||
#endif
|
||||
for (; mi < M; mi++) {
|
||||
std::memcpy(dst_col + (size_t)mi * Cin, src, 8 * sizeof(float));
|
||||
src += 8;
|
||||
}
|
||||
} else {
|
||||
const size_t copy_bytes = (size_t)n_valid_in * sizeof(float);
|
||||
for (; mi < M; mi++) {
|
||||
std::memcpy(dst_col + (size_t)mi * Cin, src, copy_bytes);
|
||||
src += 8;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// Row-major (M_chunk, Cout) -> NCHWc8 with fused BN + bias + activation
|
||||
// in the same pass; writes stream within one C-block.
|
||||
void scatterAndActivate(const float* src, int M, int p_start,
|
||||
int K1, int HW, float* out) const
|
||||
{
|
||||
const int Cout = mlas_packed_M_;
|
||||
const float* biasptr = bias.empty() ? nullptr : bias.ptr<float>();
|
||||
const float* scaleptr = fusedBatchNorm ? fusedScale.ptr<float>() : nullptr;
|
||||
const float* fbiasptr = fusedBatchNorm ? fusedBias.ptr<float>() : biasptr;
|
||||
const FastActivation act = fastActivation;
|
||||
// PReLU has per-channel slopes; LEAKY_RELU broadcasts a single one.
|
||||
const float* alphaptr = (act == FAST_ACTIV_PRELU) ? activParams.data() : nullptr;
|
||||
const float leaky_alpha = (act == FAST_ACTIV_LEAKY_RELU) ? activParams[0] : 0.f;
|
||||
const float clip_min = (act == FAST_ACTIV_CLIP) ? activParams[0] : -FLT_MAX;
|
||||
const float clip_max = (act == FAST_ACTIV_CLIP) ? activParams[1] : FLT_MAX;
|
||||
|
||||
cv::parallel_for_(cv::Range(0, K1), [&](const cv::Range& r) {
|
||||
for (int c1 = r.start; c1 < r.end; c1++) {
|
||||
const int c_base = c1 * 8;
|
||||
const int n_valid = std::min(8, Cout - c_base);
|
||||
float* dst = out + (size_t)c1 * (size_t)HW * 8
|
||||
+ (size_t)p_start * 8;
|
||||
const float* row = src + c_base;
|
||||
|
||||
CV_DECL_ALIGNED(32) float scalebuf[8] = {0};
|
||||
CV_DECL_ALIGNED(32) float biasbuf[8] = {0};
|
||||
CV_DECL_ALIGNED(32) float alphabuf[8] = {0};
|
||||
for (int c0 = 0; c0 < n_valid; c0++) {
|
||||
scalebuf[c0] = scaleptr ? scaleptr[c_base + c0] : 1.f;
|
||||
biasbuf[c0] = fbiasptr ? fbiasptr[c_base + c0] : 0.f;
|
||||
alphabuf[c0] = alphaptr ? alphaptr[c_base + c0] : leaky_alpha;
|
||||
}
|
||||
|
||||
#if CV_SIMD256 && defined(__AVX2__)
|
||||
if (n_valid == 8) {
|
||||
const __m256 vscale = _mm256_loadu_ps(scalebuf);
|
||||
const __m256 vbias = _mm256_loadu_ps(biasbuf);
|
||||
const __m256 valpha = _mm256_loadu_ps(alphabuf);
|
||||
const __m256 vzero = _mm256_setzero_ps();
|
||||
const __m256 vlo = _mm256_set1_ps(clip_min);
|
||||
const __m256 vhi = _mm256_set1_ps(clip_max);
|
||||
for (int mi = 0; mi < M; mi++) {
|
||||
__m256 v = _mm256_loadu_ps(row + (size_t)mi * Cout);
|
||||
if (scaleptr) {
|
||||
v = _mm256_fmadd_ps(v, vscale, vbias);
|
||||
} else if (fbiasptr) {
|
||||
v = _mm256_add_ps(v, vbias);
|
||||
}
|
||||
if (act == FAST_ACTIV_RELU) {
|
||||
v = _mm256_max_ps(v, vzero);
|
||||
} else if (act == FAST_ACTIV_LEAKY_RELU || act == FAST_ACTIV_PRELU) {
|
||||
__m256 neg = _mm256_mul_ps(v, valpha);
|
||||
__m256 mask = _mm256_cmp_ps(v, vzero, _CMP_GE_OQ);
|
||||
v = _mm256_blendv_ps(neg, v, mask);
|
||||
} else if (act == FAST_ACTIV_CLIP) {
|
||||
v = _mm256_min_ps(_mm256_max_ps(v, vlo), vhi);
|
||||
}
|
||||
_mm256_storeu_ps(dst + (size_t)mi * 8, v);
|
||||
}
|
||||
} else
|
||||
#endif
|
||||
{
|
||||
for (int mi = 0; mi < M; mi++) {
|
||||
for (int c0 = 0; c0 < 8; c0++) {
|
||||
float v = (c0 < n_valid) ? row[(size_t)mi * Cout + c0] : 0.f;
|
||||
if (scaleptr) v = v * scalebuf[c0] + biasbuf[c0];
|
||||
else if (fbiasptr) v = v + biasbuf[c0];
|
||||
if (act == FAST_ACTIV_RELU) v = v > 0.f ? v : 0.f;
|
||||
else if (act == FAST_ACTIV_LEAKY_RELU || act == FAST_ACTIV_PRELU)
|
||||
v = v > 0.f ? v : v * alphabuf[c0];
|
||||
else if (act == FAST_ACTIV_CLIP)
|
||||
v = std::min(std::max(v, clip_min), clip_max);
|
||||
dst[(size_t)mi * 8 + c0] = v;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
if (activationFunc) {
|
||||
float* dst = out;
|
||||
int total = K1 * HW * 8;
|
||||
activationFunc(dst, dst, total, activParams.data());
|
||||
}
|
||||
}
|
||||
|
||||
// Chunked 1x1 SGEMM-as-conv: gather NCHWc8 -> SGEMM -> scatter+activate
|
||||
// per spatial chunk. Scratch is ~CHUNK*(Cin+Cout)*4 bytes.
|
||||
void forwardMlas1x1(const Mat& inp, Mat& out)
|
||||
{
|
||||
const MatShape& inpshape = inp.shape();
|
||||
const MatShape& outshape = out.shape();
|
||||
const int Cout = mlas_packed_M_;
|
||||
const int Cin = mlas_packed_K_;
|
||||
const int ndims = inpshape.dims;
|
||||
const int N = inpshape[0];
|
||||
const int H = ndims >= 5 ? inpshape[ndims - 3] : 1;
|
||||
const int W = inpshape[ndims - 2];
|
||||
const int HW = H * W;
|
||||
const int C1_in = (Cin + 7) / 8;
|
||||
const int C1_out = (Cout + 7) / 8;
|
||||
|
||||
CV_Assert(inpshape.back() == 8 && outshape.back() == 8);
|
||||
CV_Assert(C1_in == inpshape[1]);
|
||||
|
||||
const int CHUNK = std::min(HW, 1024);
|
||||
const float* inpdata = inp.ptr<float>();
|
||||
float* outdata = out.ptr<float>();
|
||||
|
||||
// Per-layer scratch, grown lazily.
|
||||
if ((int)scratch_A_.total() < CHUNK * Cin)
|
||||
scratch_A_.create(1, CHUNK * Cin, CV_32F);
|
||||
if ((int)scratch_C_.total() < CHUNK * Cout)
|
||||
scratch_C_.create(1, CHUNK * Cout, CV_32F);
|
||||
|
||||
const size_t in_batch_stride = (size_t)C1_in * (size_t)HW * 8;
|
||||
const size_t out_batch_stride = (size_t)C1_out * (size_t)HW * 8;
|
||||
|
||||
for (int n = 0; n < N; n++) {
|
||||
const float* inp_n = inpdata + (size_t)n * in_batch_stride;
|
||||
float* out_n = outdata + (size_t)n * out_batch_stride;
|
||||
|
||||
for (int p_start = 0; p_start < HW; p_start += CHUNK) {
|
||||
int p_end = std::min(p_start + CHUNK, HW);
|
||||
int M = p_end - p_start;
|
||||
|
||||
gatherNCHWcToNHWC(inp_n, C1_in, HW, p_start, p_end,
|
||||
scratch_A_.ptr<float>(), Cin);
|
||||
|
||||
// C = A @ W^T (W is (Cout, Cin), packed with trans_b=true).
|
||||
bool ok = mlasSgemmPacked(false, true, M, Cout, Cin,
|
||||
1.0f,
|
||||
scratch_A_.ptr<float>(), Cin,
|
||||
mlas_packed_B_.data,
|
||||
0.0f,
|
||||
scratch_C_.ptr<float>(), Cout);
|
||||
CV_Assert(ok);
|
||||
|
||||
scatterAndActivate(scratch_C_.ptr<float>(), M, p_start,
|
||||
C1_out, HW, out_n);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<int> emptyKernelShape;
|
||||
Ptr<Layer> activ, batchNorm;
|
||||
Mat weights, bias, fusedScale, fusedBias;
|
||||
@@ -412,6 +667,13 @@ public:
|
||||
ActivationFunc activationFunc;
|
||||
std::vector<float> activParams;
|
||||
bool addResidual;
|
||||
|
||||
// MLAS 1x1 fast path: prepacked weight + per-layer scratch.
|
||||
Mat mlas_packed_B_;
|
||||
Mat scratch_A_;
|
||||
Mat scratch_C_;
|
||||
int mlas_packed_M_ = 0; // == Cout
|
||||
int mlas_packed_K_ = 0; // == Cin
|
||||
};
|
||||
|
||||
Ptr<Conv2Layer> Conv2Layer::create(const LayerParams& params)
|
||||
|
||||
@@ -63,6 +63,13 @@ static void deconvBlock32f(const void* inp__, const void* /*residual*/,
|
||||
const float* wdata = (const float*)weights__;
|
||||
const float* bias = bias__;
|
||||
|
||||
#if (CV_SIMD || CV_SIMD_SCALABLE)
|
||||
// SIMD path is safe when ngroups==1 or Kg%C0==0; repackDeconvWeights()
|
||||
// zero-fills padded lanes.
|
||||
const bool simd_ok =
|
||||
((ngroups == 1) || (Kg % C0 == 0)) && (C0 % VTraits<v_float32>::vlanes() == 0);
|
||||
#endif
|
||||
|
||||
const int NK1 = N * K1;
|
||||
parallel_for_(Range(0, NK1), [&](const Range& range) {
|
||||
for (int nk1 = range.start; nk1 < range.end; nk1++) {
|
||||
@@ -74,22 +81,30 @@ static void deconvBlock32f(const void* inp__, const void* /*residual*/,
|
||||
if (currK0 <= 0) continue;
|
||||
|
||||
float* out_k1 = (float*)out__ + ((int64_t)n * K1 + k1) * ospatial * C0;
|
||||
|
||||
for (int opos = 0; opos < ospatial; opos++) {
|
||||
float* p = out_k1 + opos * C0;
|
||||
if (bias) {
|
||||
for (int k0 = 0; k0 < currK0; k0++) p[k0] = bias[k_base + k0];
|
||||
} else {
|
||||
for (int k0 = 0; k0 < currK0; k0++) p[k0] = 0.f;
|
||||
}
|
||||
for (int k0 = currK0; k0 < C0; k0++) p[k0] = 0.f;
|
||||
}
|
||||
|
||||
const float* inp_n = (const float*)inp__ + (int64_t)n * C1 * ispatial * C0;
|
||||
|
||||
#if (CV_SIMD || CV_SIMD_SCALABLE)
|
||||
// Precompute the shared (g, kblk, c1_abs_base) for the SIMD path.
|
||||
int simd_g = 0, simd_kblk = 0, simd_c1_abs_base = 0;
|
||||
if (simd_ok) {
|
||||
simd_g = k_base / Kg;
|
||||
simd_kblk = (k_base - simd_g * Kg) / K0;
|
||||
const int c_start = simd_g * Cg;
|
||||
simd_c1_abs_base = c_start / C0;
|
||||
}
|
||||
#endif
|
||||
|
||||
for (int opos_flat = 0; opos_flat < ospatial; opos_flat++) {
|
||||
float* out_ptr = out_k1 + opos_flat * C0;
|
||||
|
||||
// Bias init: write currK0 valid lanes + zero-pad the rest.
|
||||
if (bias) {
|
||||
for (int k0 = 0; k0 < currK0; k0++) out_ptr[k0] = bias[k_base + k0];
|
||||
} else {
|
||||
for (int k0 = 0; k0 < currK0; k0++) out_ptr[k0] = 0.f;
|
||||
}
|
||||
for (int k0 = currK0; k0 < C0; k0++) out_ptr[k0] = 0.f;
|
||||
|
||||
int ocoords[MAX_DIMS];
|
||||
{
|
||||
int tmp = opos_flat;
|
||||
@@ -100,6 +115,64 @@ static void deconvBlock32f(const void* inp__, const void* /*residual*/,
|
||||
}
|
||||
}
|
||||
|
||||
#if (CV_SIMD || CV_SIMD_SCALABLE)
|
||||
if (simd_ok) {
|
||||
const int VLANES = VTraits<v_float32>::vlanes();
|
||||
const int v_per_block = C0 / VLANES;
|
||||
const bool two_vecs = (v_per_block == 2);
|
||||
if (v_per_block == 1 || two_vecs) {
|
||||
v_float32 acc0 = vx_load(out_ptr);
|
||||
v_float32 acc1 = two_vecs ? vx_load(out_ptr + VLANES)
|
||||
: vx_setzero_f32();
|
||||
|
||||
for (int ks = 0; ks < ksize; ks++) {
|
||||
bool valid = true;
|
||||
int ipos_flat = 0;
|
||||
for (int i = 0; i < sdims; i++) {
|
||||
int di = MAX_DIMS - sdims + i;
|
||||
int raw = ocoords[di] + cs.pads[di]
|
||||
- kcoords_tab[ks][di] * cs.dilations[di];
|
||||
if (raw < 0 || raw % cs.strides[di] != 0) {
|
||||
valid = false; break;
|
||||
}
|
||||
int ic = raw / cs.strides[di];
|
||||
if (ic >= iDims[di]) { valid = false; break; }
|
||||
ipos_flat = ipos_flat * iDims[di] + ic;
|
||||
}
|
||||
if (!valid) continue;
|
||||
|
||||
const float* w_base = wdata +
|
||||
((int64_t)(simd_g * Kblk + simd_kblk) * ksize + ks)
|
||||
* C1Max * C0 * K0;
|
||||
|
||||
for (int c1p = 0; c1p < C1Max; c1p++) {
|
||||
const int c1_abs = simd_c1_abs_base + c1p;
|
||||
if (c1_abs >= C1) break;
|
||||
|
||||
const float* inp_ptr = inp_n +
|
||||
(int64_t)(c1_abs * ispatial + ipos_flat) * C0;
|
||||
const float* w_c1p = w_base + (int64_t)c1p * C0 * K0;
|
||||
|
||||
for (int c0 = 0; c0 < C0; c0++) {
|
||||
v_float32 inp_v = vx_setall_f32(inp_ptr[c0]);
|
||||
v_float32 w0 = vx_load(w_c1p + c0 * K0);
|
||||
acc0 = v_fma(inp_v, w0, acc0);
|
||||
if (two_vecs) {
|
||||
v_float32 w1 = vx_load(w_c1p + c0 * K0 + VLANES);
|
||||
acc1 = v_fma(inp_v, w1, acc1);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
vx_store(out_ptr, acc0);
|
||||
if (two_vecs) vx_store(out_ptr + VLANES, acc1);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
#endif
|
||||
|
||||
// Scalar fallback.
|
||||
for (int ks = 0; ks < ksize; ks++) {
|
||||
bool valid = true;
|
||||
int ipos_flat = 0;
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -288,6 +288,7 @@ public:
|
||||
const int na = shape_A[shape_A.size() - 1];
|
||||
const int ma = shape_A[shape_A.size() - 2];
|
||||
const int N = shape_Y[shape_Y.size() - 1];
|
||||
const int M = shape_Y[shape_Y.size() - 2];
|
||||
const int K = trans_a ? ma : na;
|
||||
const Mat& Bmat = blobs[0];
|
||||
const int ldb = Bmat.size[Bmat.dims - 1];
|
||||
@@ -297,6 +298,7 @@ public:
|
||||
if (mlasSgemmPackB(trans_a, trans_b, N, K,
|
||||
Bmat.ptr<const float>(), ldb,
|
||||
packed_B_mlas.data)) {
|
||||
packed_B_mlas_M = M;
|
||||
packed_B_mlas_N = N;
|
||||
packed_B_mlas_K = K;
|
||||
} else {
|
||||
@@ -362,7 +364,18 @@ public:
|
||||
}
|
||||
int step = M * N;
|
||||
CV_CheckEQ(broadcast_C.size(), static_cast<size_t>(step), "DNN/Gemm: C is not broadcast properly");
|
||||
std::memcpy(ptr_y, broadcast_C.data(), step * sizeof(float));
|
||||
const float* src = broadcast_C.data();
|
||||
constexpr size_t PAR_THRESHOLD = 1u << 14;
|
||||
if ((size_t)step < PAR_THRESHOLD) {
|
||||
std::memcpy(ptr_y, src, (size_t)step * sizeof(float));
|
||||
} else {
|
||||
int nthreads = std::max(1, getNumThreads());
|
||||
parallel_for_(Range(0, nthreads), [=](const Range& r) {
|
||||
size_t st = (size_t)r.start * (size_t)step / (size_t)nthreads;
|
||||
size_t en = (size_t)r.end * (size_t)step / (size_t)nthreads;
|
||||
std::memcpy(ptr_y + st, src + st, (en - st) * sizeof(float));
|
||||
}, nthreads);
|
||||
}
|
||||
} else {
|
||||
// ND output: tile the (1D / scalar) bias across all `rows`
|
||||
// rows, scaled by beta. The rewriter restricts bias to scalar
|
||||
@@ -371,19 +384,55 @@ public:
|
||||
const float* c = C.ptr<const float>();
|
||||
if (C.total() == 1) {
|
||||
float val = beta * (*c);
|
||||
std::fill_n(ptr_y, (size_t)rows * (size_t)N, val);
|
||||
size_t total = (size_t)rows * (size_t)N;
|
||||
constexpr size_t PAR_THRESHOLD = 1u << 14;
|
||||
if (total < PAR_THRESHOLD) {
|
||||
std::fill_n(ptr_y, total, val);
|
||||
} else {
|
||||
int nthreads = std::max(1, getNumThreads());
|
||||
parallel_for_(Range(0, nthreads), [=](const Range& r) {
|
||||
size_t st = (size_t)r.start * total / (size_t)nthreads;
|
||||
size_t en = (size_t)r.end * total / (size_t)nthreads;
|
||||
std::fill_n(ptr_y + st, en - st, val);
|
||||
}, nthreads);
|
||||
}
|
||||
} else {
|
||||
CV_CheckEQ((int)C.total(), N, "DNN/Gemm: bias must be scalar or length-N in flatten_a=false mode");
|
||||
// Compute first row, then broadcast.
|
||||
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));
|
||||
if (rows > 1) {
|
||||
size_t row_bytes = (size_t)N * sizeof(float);
|
||||
constexpr int PAR_THRESHOLD_ROWS = 32;
|
||||
if (rows <= PAR_THRESHOLD_ROWS) {
|
||||
for (int i = 1; i < rows; i++) {
|
||||
std::memcpy(ptr_y + (size_t)i * N, ptr_y, row_bytes);
|
||||
}
|
||||
} else {
|
||||
int nthreads = std::max(1, getNumThreads());
|
||||
double nstripes = std::min((double)(rows - 1), (double)nthreads);
|
||||
parallel_for_(Range(1, rows), [=](const Range& r) {
|
||||
for (int i = r.start; i < r.end; i++) {
|
||||
std::memcpy(ptr_y + (size_t)i * N, ptr_y, row_bytes);
|
||||
}
|
||||
}, nstripes);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else { // initialization
|
||||
float *ptr_y = Y.ptr<float>();
|
||||
size_t total = Y.total();
|
||||
std::memset(ptr_y, 0, total * sizeof(float));
|
||||
constexpr size_t PAR_THRESHOLD = 1u << 14;
|
||||
if (total < PAR_THRESHOLD) {
|
||||
std::memset(ptr_y, 0, total * sizeof(float));
|
||||
} else {
|
||||
int nthreads = std::max(1, getNumThreads());
|
||||
parallel_for_(Range(0, nthreads), [=](const Range& r) {
|
||||
size_t st = (size_t)r.start * total / (size_t)nthreads;
|
||||
size_t en = (size_t)r.end * total / (size_t)nthreads;
|
||||
std::memset(ptr_y + st, 0, (en - st) * sizeof(float));
|
||||
}, nthreads);
|
||||
}
|
||||
}
|
||||
|
||||
if (constB(mode)) {
|
||||
@@ -573,6 +622,7 @@ private:
|
||||
std::vector<float> thin_packed_B;
|
||||
#ifdef HAVE_MLAS
|
||||
cv::Mat packed_B_mlas;
|
||||
int packed_B_mlas_M = 0;
|
||||
int packed_B_mlas_N = 0;
|
||||
int packed_B_mlas_K = 0;
|
||||
#endif
|
||||
|
||||
@@ -226,10 +226,14 @@ private:
|
||||
class ParallelSlice : public cv::ParallelLoopBody
|
||||
{
|
||||
public:
|
||||
// parallel_axis: dim to split across threads. Leading dims (0..parallel_axis-1) are
|
||||
// collapsed into a constant offset since they have size 1 in the output.
|
||||
ParallelSlice(const Mat& inp, Mat& out,
|
||||
const std::vector<Range>& ranges,
|
||||
const std::vector<int>& steps)
|
||||
: inp_(inp), out_(out), ranges_(ranges), steps_(steps)
|
||||
const std::vector<int>& steps,
|
||||
int parallel_axis)
|
||||
: inp_(inp), out_(out), ranges_(ranges), steps_(steps),
|
||||
parallel_axis_(parallel_axis)
|
||||
{
|
||||
dims_ = inp.dims;
|
||||
es_ = inp.elemSize();
|
||||
@@ -240,25 +244,40 @@ private:
|
||||
inp_strides_[i] = inp.step.p[i];
|
||||
out_strides_[i] = out.step.p[i];
|
||||
}
|
||||
|
||||
preamble_src_ = 0;
|
||||
for (int d = 0; d < parallel_axis_; d++) {
|
||||
preamble_src_ += (size_t)ranges_[d].start * inp_strides_[d];
|
||||
}
|
||||
}
|
||||
|
||||
void operator()(const Range& range) const CV_OVERRIDE
|
||||
{
|
||||
int b = ranges_[0].start;
|
||||
int s = steps_[0];
|
||||
const int axis = parallel_axis_;
|
||||
const int b = ranges_[axis].start;
|
||||
const int s = steps_[axis];
|
||||
|
||||
const uchar* src_base = inp_.ptr();
|
||||
const uchar* src_base = inp_.ptr() + preamble_src_;
|
||||
uchar* dst_base = out_.ptr();
|
||||
|
||||
if (s == 1 && axis < dims_ - 1 && is_fully_contiguous(axis)) {
|
||||
const size_t unit_bytes = inp_strides_[axis];
|
||||
const size_t count = (size_t)(range.end - range.start);
|
||||
const size_t src_off = (size_t)(b + range.start) * unit_bytes;
|
||||
const size_t dst_off = (size_t)range.start * out_strides_[axis];
|
||||
std::memcpy(dst_base + dst_off, src_base + src_off, count * unit_bytes);
|
||||
return;
|
||||
}
|
||||
|
||||
for (int i = range.start; i < range.end; ++i)
|
||||
{
|
||||
int k = b + i * s;
|
||||
size_t src_offset = k * inp_strides_[0];
|
||||
size_t dst_offset = i * out_strides_[0];
|
||||
if (dims_ == 1)
|
||||
size_t src_offset = (size_t)k * inp_strides_[axis];
|
||||
size_t dst_offset = (size_t)i * out_strides_[axis];
|
||||
if (axis == dims_ - 1)
|
||||
std::memcpy(dst_base + dst_offset, src_base + src_offset, es_);
|
||||
else
|
||||
recursive_copy(1, src_base + src_offset, dst_base + dst_offset);
|
||||
recursive_copy(axis + 1, src_base + src_offset, dst_base + dst_offset);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -364,14 +383,26 @@ private:
|
||||
size_t es_;
|
||||
std::vector<size_t> inp_strides_;
|
||||
std::vector<size_t> out_strides_;
|
||||
int parallel_axis_;
|
||||
size_t preamble_src_;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
void run_parallel(const Mat& inp, Mat& out, const std::vector<Range>& ranges, const std::vector<int>& steps)
|
||||
{
|
||||
ParallelSlice<T> body(inp, out, ranges, steps);
|
||||
int dim0_size = out.size[0];
|
||||
parallel_for_(Range(0, dim0_size), body);
|
||||
int dims = inp.dims;
|
||||
// Skip leading size-1 dims; otherwise the whole slice is one stripe.
|
||||
int parallel_axis = 0;
|
||||
while (parallel_axis < dims - 1 && out.size[parallel_axis] == 1) {
|
||||
parallel_axis++;
|
||||
}
|
||||
int parallel_size = out.size[parallel_axis];
|
||||
ParallelSlice<T> body(inp, out, ranges, steps, parallel_axis);
|
||||
// One stripe per thread; per-element body is tiny so dynamic chunking
|
||||
// would burn cycles on task sync.
|
||||
int nthreads = std::max(1, getNumThreads());
|
||||
double nstripes = std::min((double)parallel_size, (double)nthreads);
|
||||
parallel_for_(Range(0, parallel_size), body, nstripes);
|
||||
}
|
||||
|
||||
void forward(InputArrayOfArrays inputs_arr,
|
||||
|
||||
@@ -311,11 +311,13 @@ void transformLayout(const Mat& inp, Mat& out,
|
||||
constexpr size_t min_elems_per_chunk = 1 << 14;
|
||||
int nblocks = int((total + min_elems_per_chunk/2) / min_elems_per_chunk);
|
||||
int nthreads = std::max(1, getNumThreads());
|
||||
nblocks = clamp(nblocks, 1, std::max(N*C1, 1) * 4);
|
||||
nblocks = clamp(nblocks, 1, std::max(N*C1, 1) * 16);
|
||||
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)
|
||||
int total_chunks = N * C1 * nblocks;
|
||||
double nstripes = std::min((double)total_chunks, (double)nthreads);
|
||||
parallel_for_(Range(0, total_chunks), [&](const Range& range)
|
||||
{
|
||||
int dchunk = 1u;
|
||||
bool interleave = inplayout == DATA_LAYOUT_NCHW;
|
||||
@@ -341,7 +343,7 @@ void transformLayout(const Mat& inp, Mat& out,
|
||||
|
||||
kernel(inptr + inpofs, outptr + outofs, C, planesize, nc, nzc, dlen);
|
||||
}
|
||||
});
|
||||
}, nstripes);
|
||||
}
|
||||
|
||||
class TransformLayoutLayerImpl : public TransformLayoutLayer
|
||||
|
||||
@@ -2738,7 +2738,7 @@ void Net::Impl::collectOrtProfileData() const
|
||||
for (auto& kv : by_name) {
|
||||
const String name = kv.first;
|
||||
const String type = kv.second.first;
|
||||
const double ms_per_run = kv.second.second / runs;
|
||||
const double ms_per_run = kv.second.second / std::max(1, runs);
|
||||
ort_profile_data.emplace_back(name, type, ms_per_run);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user