Muse glimmer / dflash support

ggml/src/ggml-openvino/ggml-decoder.cpp

@@ -35,6 +35,24 @@
 #include <unordered_map>
 #include <vector>
 
+static std::set<std::string> collect_graph_output_names(const ggml_cgraph * cgraph) {
+    std::set<std::string> outputs;
+    for (int node_n = 0; node_n < cgraph->n_nodes; node_n++) {
+        auto * node = cgraph->nodes[node_n];
+        if (node->flags & GGML_TENSOR_FLAG_OUTPUT) {
+            outputs.insert(node->name);
+        }
+    }
+    return outputs;
+}
+
+bool GgmlOvDecoder::has_same_graph_io(const ggml_cgraph * cgraph) const {
+    if (m_built_output_names.empty()) {
+        return true;  // no snapshot (naive / test decoders): nothing to compare against
+    }
+    return collect_graph_output_names(cgraph) == m_built_output_names;
+}
+
 GgmlOvDecoder::GgmlOvDecoder(ggml_cgraph * cgraph,
                              ModelParams & model_params,
                              ComputeParams & compute_params,
@@ -64,6 +82,8 @@ GgmlOvDecoder::GgmlOvDecoder(ggml_cgraph * cgraph,
 
     validate_cgraph();
 
+    m_built_output_names = collect_graph_output_names(cgraph);
+
     set_input_output();
     compute_node_dynamic_dims();
     compute_model_inputs();
@@ -79,6 +99,7 @@ GgmlOvDecoder::GgmlOvDecoder(ggml_cgraph * cgraph,
 
 void GgmlOvDecoder::update_io(ggml_cgraph * cgraph) {
     m_cgraph = cgraph;
+    m_built_output_names = collect_graph_output_names(cgraph);
     m_model_inputs.clear();
     m_model_outputs.clear();
     m_node_info_list.clear();
@@ -983,7 +1004,7 @@ ov::PartialShape GgmlOvDecoder::get_graph_input_shape(const ggml_tensor * op,
     } else {
         input_shape = ov::PartialShape{get_shape(input)};
     }
-    if (dynamic_dim_index != -1 && m_model_is_splitted) {
+    if (dynamic_dim_index != -1) {
         input_shape[3 - dynamic_dim_index] = -1;
     }
     if (op->op == GGML_OP_SOFT_MAX && op->src[1] != nullptr && op->src[1]->op == GGML_OP_NONE &&
@@ -1171,6 +1192,30 @@ void GgmlOvDecoder::compute_model_outputs() {
         if (cur_node->op == GGML_OP_NONE || cur_node->op == GGML_OP_VIEW || cur_node->op == GGML_OP_RESHAPE) {
             continue;
         }
+        // A node explicitly marked by ggml_set_output() must be a model output even when it is
+        // fully consumed inside the graph, because the host reads it back afterwards. The
+        // use-count analysis below cannot see that: it only asks whether anything still in the
+        // graph needs the value.
+        //
+        // Mid-graph taps break that assumption. A model that exposes its residual stream at a few
+        // layers (llm_graph_result marks each such tensor with ggml_set_output()) has tensors that
+        // are read by the host AND feed the next layer, so use-count logic dropped them, the OV
+        // model never produced them, and the host read an unwritten buffer -- whatever consumes
+        // those features saw zeros. This affects any consumer of a mid-graph tensor, not one
+        // particular feature.
+        if (cur_node->flags & GGML_TENSOR_FLAG_OUTPUT) {
+            // in-place ops publish their result through view_src, as in the use_count == 0 path
+            auto * flagged_node = cur_node;
+            if (flagged_node->op == GGML_OP_SET_ROWS && flagged_node->view_src != nullptr) {
+                flagged_node = flagged_node->view_src;
+            }
+            std::string flagged_name = get_tensor_ov_name(m_cgraph, flagged_node);
+            if (m_model_outputs.find(flagged_name) == m_model_outputs.end()) {
+                m_model_outputs[flagged_name] = flagged_node;
+                m_model_output_names.push_back(flagged_name);
+            }
+            continue;
+        }
         auto cur_node_use_count = m_cgraph->use_counts[ggml_hash_find(&m_cgraph->visited_hash_set, cur_node)];
         if (cur_node_use_count == 0) {
             // The output of in-place ops is the view_src tensor, which is updated in place. We should use the view_src name as the output name to make sure it can be correctly matched with the later ops that use the view_src.
@@ -1193,8 +1238,10 @@ void GgmlOvDecoder::compute_model_outputs() {
         }
         if (cur_node != nullptr) {
             std::string cur_node_name = get_tensor_ov_name(m_cgraph, cur_node);
+            if (m_model_outputs.find(cur_node_name) == m_model_outputs.end()) {
+                m_model_output_names.push_back(cur_node_name);
+            }
             m_model_outputs[cur_node_name] = cur_node;
-            m_model_output_names.insert(cur_node_name);
         }
     }
 }
@@ -1899,6 +1946,10 @@ void GgmlOvDecoder::compute_node_dynamic_dims() {
                     m_node_dynamic_dims[src] = 1;
                     continue;
                 }
+                if (is_inp_embd_2d(src, node)) {
+                    m_node_dynamic_dims[src] = 1;
+                    continue;
+                }
                 self(self, src);
             }
         }
@@ -1954,6 +2005,11 @@ void GgmlOvDecoder::compute_node_dynamic_dims() {
             }
             break;
         case GGML_OP_VIEW: {
+            // A VIEW that does not change the shape is a pass-through for dynamic-dim purposes.
+            if (is_same_shape(node->src[0], node)) {
+                m_node_dynamic_dims[node] = m_node_dynamic_dims[node->src[0]];
+                break;
+            }
             // Use stride-based matching: the stride of a VIEW dimension directly
             // encodes which source dimension it indexes into, so it uniquely
             // identifies the dynamic dim even when two dims share the same size.

ggml/src/ggml-openvino/ggml-decoder.h

@@ -8,10 +8,10 @@
 #include <cstdint>
 #include <cstring>
 #include <map>
+#include <set>
 #include <memory>
 #include <openvino/core/partial_shape.hpp>
 #include <optional>
-#include <set>
 #include <string>
 #include <vector>
 
@@ -250,7 +250,9 @@ class GgmlOvDecoder : public ov::frontend::ggml::GgmlDecoder {
         return m_model_weights;
     }
 
-    virtual std::set<std::string> get_model_output_names() const override { return m_model_output_names; }
+    virtual std::set<std::string> get_model_output_names() const override {
+        return std::set<std::string>(m_model_output_names.begin(), m_model_output_names.end());
+    }
 
     const std::map<std::string, ggml_tensor *> & get_model_outputs() const { return m_model_outputs; }
 
@@ -349,6 +351,16 @@ class GgmlOvDecoder : public ov::frontend::ggml::GgmlDecoder {
 
     void update_io(ggml_cgraph * cgraph);
 
+    // The cgraph this decoder currently describes. A cached decoder can be handed a DIFFERENT
+    // cgraph instance with the same graph_key, and every cached ggml_tensor* (and m_cgraph itself)
+    // then dangles, so callers must compare and refresh.
+    const ggml_cgraph * get_cgraph() const { return m_cgraph; }
+
+    // Was this decoder built for a graph with the same OUTPUT set as `cgraph`? The cache key is
+    // only (n_nodes, first/last node name), which cannot see a change in which tensors are marked
+    // as outputs, and the compiled model's Results are fixed at compile time.
+    bool has_same_graph_io(const ggml_cgraph * cgraph) const;
+
     inline static bool is_inp_tok(const ggml_tensor * tensor, const ggml_tensor * op) {
         return op->op == GGML_OP_GET_ROWS && tensor == op->src[1] && op->src[0]->op == GGML_OP_NONE;
     }
@@ -367,6 +379,16 @@ class GgmlOvDecoder : public ov::frontend::ggml::GgmlDecoder {
         return tensor->op == GGML_OP_GET_ROWS && op->op == GGML_OP_RMS_NORM;
     }
 
+    // A 2-D float graph input feeding a matmul as src[1]. Speculative decoding hands the draft the
+    // target's hidden features this way ([n_embd, n_tokens], with no GET_ROWS lookup, so none of the
+    // other seeds match it) and the token count differs between the prompt pass and the shorter
+    // draft blocks, so its token dim must stay dynamic.
+    inline static bool is_inp_embd_2d(const ggml_tensor * tensor, const ggml_tensor * op) {
+        return (tensor->flags & GGML_TENSOR_FLAG_INPUT) && tensor->op == GGML_OP_NONE &&
+               tensor->type == GGML_TYPE_F32 && tensor->ne[2] == 1 && tensor->ne[3] == 1 &&
+               op->op == GGML_OP_MUL_MAT && tensor == op->src[1];
+    }
+
     inline static bool is_inp_mask(const ggml_tensor * tensor, const ggml_tensor * op) {
         return op->op == GGML_OP_CPY || (op->op == GGML_OP_FLASH_ATTN_EXT && tensor == op->src[3]) ||
                (op->op == GGML_OP_SOFT_MAX && tensor == op->src[1]);
@@ -454,7 +476,9 @@ class GgmlOvDecoder : public ov::frontend::ggml::GgmlDecoder {
     std::map<std::string, ov::frontend::ggml::ModelExtraInputInfo> m_model_extra_inputs;
     std::map<std::string, std::shared_ptr<ov::Node>> m_model_weights;
     std::map<std::string, ggml_tensor *> m_model_outputs;
-    std::set<std::string> m_model_output_names;
+    std::vector<std::string> m_model_output_names;
+    // The output set the compiled model was built from, for has_same_graph_io().
+    std::set<std::string> m_built_output_names;
     std::vector<NodeInfo> m_node_info_list;
     std::map<ggml_tensor *, int> m_node_dynamic_dims;
 

ggml/src/ggml-openvino/ggml-openvino-extra.cpp

@@ -57,9 +57,12 @@ void ggml_openvino_device_config::init() {
         "GGML_OPENVINO_RELEASE_WEIGHTS",
         "GGML_OPENVINO_REDUCE_COMPILE_MEM",
         "GGML_OPENVINO_LOG_UNSUPPORTED_OPS",
+        "GGML_OPENVINO_ROPE_F32_ANGLE",
         "GGML_OPENVINO_LOG_SWA_LAYERS",
         "GGML_OPENVINO_REQUANT_KQUANT",
         "GGML_OPENVINO_DISABLE_KV_STATE_RELAYOUT",
+        "GGML_OPENVINO_LOG_SUPPORTS_OP",
+        "GGML_OPENVINO_DIAG_GRAPH_IO",
     };
 
     for (const char * const & env_var : env_var_names) {

ggml/src/ggml-openvino/ggml-openvino.cpp

@@ -16,6 +16,7 @@
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
+#include <map>
 #include <memory>
 #include <mutex>
 #include <openvino/core/type/element_type.hpp>
@@ -1523,16 +1524,114 @@ static ggml_openvino_op_support ggml_backend_openvino_device_supports_op_impl(gg
     return {true, ""};
 }
 
+// Report what supports_op() accepts and rejects, so a graph that silently falls back to another
+// backend can be diagnosed. Every op rejected here becomes a graph split; if nothing is rejected,
+// the whole graph runs on OpenVINO.
+//
+// Gated by GGML_OPENVINO_LOG_SUPPORTS_OP:
+//   1  one summary at exit listing every distinct op kind and verdict -- start here
+//   2  additionally log every individual call, with the tensor name and shape
+//
+// This wraps the decision function instead of instrumenting it, so all of its return paths are
+// covered -- including the several that reject without logging a reason of their own, and any
+// added later.
+//
+// The summary is reported at exit rather than on each op's first occurrence, because the earliest
+// calls happen before the log handler emits anything: with level 2 the first line to actually
+// appear already reads "call 1244". A first-occurrence line would therefore be silently dropped
+// for exactly the ops that are only ever seen early. Counting is unaffected, so the summary is the
+// reliable view; treat the level-2 stream as a sample, not a complete trace.
+namespace {
+
+struct supports_op_tally {
+    std::mutex mutex;
+    // key: "<op>:ok" / "<op>:no" -> number of calls
+    std::map<std::string, size_t> counts;
+
+    // Report totals once, at exit. Registered lazily so a run with the knob off costs nothing.
+    // atexit runs while the ggml log callback is still valid, which static destruction does not
+    // guarantee.
+    void arm_summary() {
+        static std::once_flag armed;
+        std::call_once(armed, [] {
+            std::atexit([] {
+                auto & self = instance();
+                std::lock_guard<std::mutex> lock(self.mutex);
+
+                size_t n_accept_kinds = 0;
+                size_t n_reject_kinds = 0;
+                GGML_LOG_WARN("ov-supports-op: ---- summary (distinct op kinds, with call counts) ----\n");
+                for (const auto & [key, count] : self.counts) {
+                    const bool ok = key.size() > 3 && key.compare(key.size() - 3, 3, ":ok") == 0;
+                    ok ? ++n_accept_kinds : ++n_reject_kinds;
+                    GGML_LOG_WARN("ov-supports-op:   %-8s %-24s %zu call(s)\n", ok ? "ACCEPT" : "REJECT",
+                                  key.substr(0, key.size() - 3).c_str(), count);
+                }
+                if (n_reject_kinds == 0) {
+                    GGML_LOG_WARN("ov-supports-op: %zu op kind(s) accepted, 0 rejected"
+                                  " -- every op was supported, so no split originated here\n",
+                                  n_accept_kinds);
+                } else {
+                    GGML_LOG_WARN("ov-supports-op: %zu op kind(s) accepted, %zu REJECTED"
+                                  " -- each rejected kind forces a graph split\n",
+                                  n_accept_kinds, n_reject_kinds);
+                }
+            });
+        });
+    }
+
+    static supports_op_tally & instance() {
+        static supports_op_tally tally;
+        return tally;
+    }
+};
+
+}  // namespace
+
 static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) {
     auto res = ggml_backend_openvino_device_supports_op_impl(dev, op);
-    if (!res.is_supported) {
+    const bool supported = res.is_supported;
+    if (!supported) {
         static const bool log_unsupported = ggml_openvino_getenv_int("GGML_OPENVINO_LOG_UNSUPPORTED_OPS") != 0;
         if (log_unsupported) {
             GGML_LOG_WARN("OpenVINO op unsupported: op '%s' (%s), type %s: %s\n",
                           op->name, ggml_op_name(op->op), ggml_type_name(op->type), res.reason.c_str());
         }
     }
-    return res.is_supported;
+
+    static const int log_level = ggml_openvino_getenv_int("GGML_OPENVINO_LOG_SUPPORTS_OP");
+    if (log_level <= 0) {
+        return supported;
+    }
+
+    // Name the op the way ggml does: UNARY and GLU carry their real identity in a sub-enum, so
+    // logging only "UNARY" would merge SILU with everything else.
+    std::string op_desc = ggml_op_name(op->op);
+    if (op->op == GGML_OP_UNARY) {
+        op_desc += std::string("/") + ggml_unary_op_name(ggml_get_unary_op(op));
+    } else if (op->op == GGML_OP_GLU) {
+        op_desc += std::string("/") + ggml_glu_op_name(ggml_get_glu_op(op));
+    }
+
+    // The scheduler calls supports_op() on every graph build and may do so from several threads,
+    // so the tally needs a lock. Keyed on op identity *and* verdict, because the same op can be
+    // accepted at one shape and rejected at another (e.g. a quantized 3-D src).
+    auto & tally = supports_op_tally::instance();
+    tally.arm_summary();
+
+    size_t count;
+    {
+        std::lock_guard<std::mutex> lock(tally.mutex);
+        count = ++tally.counts[op_desc + (supported ? ":ok" : ":no")];
+    }
+
+    if (log_level >= 2) {
+        GGML_LOG_WARN("ov-supports-op: %-6s %-22s name='%s' type=%s ne=[%ld,%ld,%ld,%ld] (call %zu)\n",
+                      supported ? "ACCEPT" : "REJECT", op_desc.c_str(), op->name, ggml_type_name(op->type),
+                      (long) op->ne[0], (long) op->ne[1], (long) op->ne[2], (long) op->ne[3], count);
+    }
+
+    return supported;
 }
 
 static bool ggml_backend_openvino_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) {

ggml/src/ggml-openvino/openvino/op/argmax.cpp

@@ -0,0 +1,51 @@
+#include "../node_context.h"
+#include "../op_table.h"
+#include "../utils.h"
+#include "ggml.h"
+
+#include <openvino/frontend/exception.hpp>
+#include <openvino/op/concat.hpp>
+#include <openvino/op/constant.hpp>
+#include <openvino/op/reshape.hpp>
+#include <openvino/op/topk.hpp>
+
+namespace ov {
+namespace frontend {
+namespace ggml {
+namespace op {
+
+// GGML_OP_ARGMAX: src0 is a matrix [ne0, ne1]; the result is a 1-D I32 tensor of ne1 entries
+// holding, for each row, the index of the maximum along ne0. ggml shapes arrive reversed here, so
+// ne0 is the last OV axis and ne1 is axis 2. TopK(k=1, axis=last, MAX) yields exactly that index,
+// then reshape to the [1,1,1,ne1] layout the decoder expects for a 1-D ggml output.
+//
+// Note ggml's argmax returns the FIRST maximum on a tie, while TopK's tie-breaking is not specified
+// to match. Ties on real logits are vanishingly rare and none were observed, but the two are not
+// proven equivalent.
+OutputVector translate_argmax(const NodeContext & context) {
+    num_inputs_check(context, 1, 1);
+
+    auto input = process_view_input_new(context, 0);
+
+    auto k = ov::op::v0::Constant::create(ov::element::i64, {}, {1});
+    auto topk = std::make_shared<ov::op::v11::TopK>(input,
+                                                   k,
+                                                   3,
+                                                   ov::op::v11::TopK::Mode::MAX,
+                                                   ov::op::v11::TopK::SortType::SORT_VALUES,
+                                                   context.get_output_type(),
+                                                   false);
+
+    // ne1 is the dynamic row count, so build the target shape at runtime.
+    auto leading = ov::op::v0::Constant::create(ov::element::i64, {3}, {1, 1, 1});
+    auto rows = get_dimensions(input.get_node_shared_ptr(), {2});
+    auto target_shape = std::make_shared<ov::op::v0::Concat>(OutputVector{leading, rows}, 0);
+    auto res = std::make_shared<ov::op::v1::Reshape>(topk->output(1), target_shape, false);
+
+    return rename_outputs_with_suffix({res}, context.get_name());
+}
+
+}  // namespace op
+}  // namespace ggml
+}  // namespace frontend
+}  // namespace ov

ggml/src/ggml-openvino/openvino/op_table.cpp

@@ -45,6 +45,7 @@ std::unordered_map<std::string, CreatorFunction> get_supported_ops() {
         {"GGML_OP_SQR",             op::translate_sqr                              },
         {"GGML_OP_SQRT",            op::translate_sqrt                             },
         {"GGML_OP_SOFT_MAX",        op::translate_soft_max                         },
+        {"GGML_OP_ARGMAX",          op::translate_argmax                           },
         {"GGML_OP_ARGSORT",         op::translate_argsort                          },
         {"GGML_OP_SUB",             op::translate_1to1_match_2_inputs<v1::Subtract>},
         {"GGML_OP_TRANSPOSE",       op::translate_transpose                        },

ggml/src/ggml-openvino/openvino/op_table.h

@@ -42,6 +42,7 @@ GGML_OP_CONVERTER(translate_glu_geglu);
 GGML_OP_CONVERTER(translate_glu_geglu_quick);
 GGML_OP_CONVERTER(translate_set_rows);
 GGML_OP_CONVERTER(translate_cpy);
+GGML_OP_CONVERTER(translate_argmax);
 GGML_OP_CONVERTER(translate_argsort);
 GGML_OP_CONVERTER(translate_flash_attn_ext);
 GGML_OP_CONVERTER(translate_clamp);