extract some common fusion logic

Author: jeffbolznv

ggml/src/ggml-impl.h

@@ -589,6 +589,28 @@ static inline ggml_bf16_t ggml_compute_fp32_to_bf16(float s) {
 #define GGML_FP32_TO_BF16(x) ggml_compute_fp32_to_bf16(x)
 #define GGML_BF16_TO_FP32(x) ggml_compute_bf16_to_fp32(x)
 
+// return true if the node's results are only used by N other nodes
+// and can be fused into their calculations.
+static inline bool ggml_can_fuse_node(const struct ggml_tensor * node, int32_t N) {
+    // check the use count against how many we're replacing
+    if (node->use_count != N) {
+        return false;
+    }
+
+    // if node is a view, some other node might be using the intermediate result
+    // via the view source.
+    if (node->view_src) {
+        return false;
+    }
+
+    // If the user requested output for the node, can't fuse
+    if (node->flags & GGML_TENSOR_FLAG_OUTPUT) {
+        return false;
+    }
+
+    return true;
+}
+
 #ifdef __cplusplus
 }
 #endif

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -9698,15 +9698,14 @@ static bool ggml_vk_is_empty(ggml_tensor * node) {
 }
 
 // Returns true if nodes [i, i+1] are fusable RMS_NORM + MUL.
-bool ggml_can_fuse_rms_norm_mul(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, int i) {
+static bool ggml_can_fuse_rms_norm_mul(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, int i) {
     ggml_tensor *norm = cgraph->nodes[i];
 
-    if (norm->op != GGML_OP_RMS_NORM || norm->use_count != 1) {
+    if (norm->op != GGML_OP_RMS_NORM) {
         return false;
     }
-    // if norm is a view, some other node might be using the intermediate result
-    // view the view source.
-    if (norm->view_src) {
+
+    if (!ggml_can_fuse_node(norm, 1)) {
         return false;
     }
 
@@ -9721,7 +9720,6 @@ bool ggml_can_fuse_rms_norm_mul(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
     // Since norm is the first operand of mul, it must be the same shape
     GGML_ASSERT(ggml_are_same_shape(mul, norm));
 
-    // XXX TODO: Do we need a way to indicate that the user doesn't need the intermediate result?
     return true;
 }