char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Zijun YuOpenVINO: Qwen3.5, memory optimization, and test-recurrent-state-rollback (#26952)aee56b3ab

char/aliceai-t5-moe
2.2 KiB58 linesraw
1#include "../node_context.h"
2#include "../op_table.h"
3#include "../utils.h"
4
5#include <openvino/op/constant.hpp>
6#include <openvino/op/equal.hpp>
7#include <openvino/op/multiply.hpp>
8#include <openvino/op/range.hpp>
9#include <openvino/op/reshape.hpp>
10#include <openvino/op/select.hpp>
11
12namespace ov {
13namespace frontend {
14namespace ggml {
15namespace op {
16
17// GGML DIAG takes a 1D vector (ne0, 1, ne2, ne3) and produces a diagonal matrix
18// of shape (ne0, ne0, ne2, ne3).
19// In OV layout (ggml [ne0, ne1, ne2, ne3] → OV [ne3, ne2, ne1, ne0]):
20//   input:  [ne3, ne2, 1, ne0]
21//   output: [ne3, ne2, ne0, ne0]
22// The diagonal: output[..., i, j] = input[..., 0, j] if i == j, else 0.
23OutputVector translate_diag(const NodeContext & context) {
24    num_inputs_check(context, 1, 1);
25
26    auto x = context.get_input(0);  // OV shape: [ne3, ne2, 1, ne0]
27
28    auto out_shape = context.get_output_shape().to_shape();
29    int64_t n = static_cast<int64_t>(out_shape[3]);  // ne0
30
31    // Build index range [0, 1, ..., n-1]
32    auto start = ov::op::v0::Constant::create(ov::element::i64, {}, {int64_t(0)});
33    auto stop  = ov::op::v0::Constant::create(ov::element::i64, {}, {n});
34    auto step  = ov::op::v0::Constant::create(ov::element::i64, {}, {int64_t(1)});
35    auto range = std::make_shared<ov::op::v4::Range>(start, stop, step, ov::element::i64);
36
37    // col_idx shape [1, 1, 1, n]
38    auto col_shape = ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{1, 1, 1, n});
39    auto col_idx   = std::make_shared<ov::op::v1::Reshape>(range, col_shape, false);
40
41    // row_idx shape [1, 1, n, 1]
42    auto row_shape = ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{1, 1, n, 1});
43    auto row_idx   = std::make_shared<ov::op::v1::Reshape>(range, row_shape, false);
44
45    // mask: true where col == row (diagonal)
46    auto mask = std::make_shared<ov::op::v1::Equal>(col_idx, row_idx);
47
48    // Broadcast input from [ne3, ne2, 1, ne0] to [ne3, ne2, ne0, ne0] via select
49    auto zero = ov::op::v0::Constant::create(ov::element::f32, {}, {0.0f});
50    auto res  = std::make_shared<ov::op::v1::Select>(mask, x, zero);
51
52    return rename_outputs_with_suffix({res}, context.get_name());
53}
54
55}  // namespace op
56}  // namespace ggml
57}  // namespace frontend
58}  // namespace ov