1
0
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:
Abhishek Gola
2026-05-27 12:02:27 +05:30
committed by GitHub
parent 2405556989
commit bf0cf34963
11 changed files with 2049 additions and 41 deletions
+514 -6
View File
@@ -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 &params) {
return makePtr<AttentionLayerImpl>(params);
}
}} // cv::dnn
+53 -1
View File
@@ -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)
+262
View File
@@ -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
+55 -5
View File
@@ -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
+43 -12
View File
@@ -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
+1 -1
View File
@@ -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);
}
}