Add --override-tensors option to llama-bench

Author: 4onen

examples/llama-bench/llama-bench.cpp

@@ -181,6 +181,7 @@ struct cmd_params {
     int                              reps;
     ggml_sched_priority              prio;
     int                              delay;
+    std::vector<llama_model_tensor_buft_override> tensor_buft_overrides;
     bool                             verbose;
     bool                             progress;
     output_formats                   output_format;
@@ -213,6 +214,7 @@ static const cmd_params cmd_params_defaults = {
     /* reps                 */ 5,
     /* prio                 */ GGML_SCHED_PRIO_NORMAL,
     /* delay                */ 0,
+    /* tensor_buft_overrides*/ {},
     /* verbose              */ false,
     /* progress             */ false,
     /* output_format        */ MARKDOWN,
@@ -268,6 +270,7 @@ static void print_usage(int /* argc */, char ** argv) {
     printf("  -r, --repetitions <n>                     (default: %d)\n", cmd_params_defaults.reps);
     printf("  --prio <0|1|2|3>                          (default: %d)\n", cmd_params_defaults.prio);
     printf("  --delay <0...N> (seconds)                 (default: %d)\n", cmd_params_defaults.delay);
+    printf("  -ot --override-tensors <tensor name pattern>=<buffer type>,... (default:disabled)\n");
     printf("  -o, --output <csv|json|jsonl|md|sql>      (default: %s)\n",
            output_format_str(cmd_params_defaults.output_format));
     printf("  -oe, --output-err <csv|json|jsonl|md|sql> (default: %s)\n",
@@ -575,6 +578,56 @@ static cmd_params parse_cmd_params(int argc, char ** argv) {
                 break;
             }
             params.delay = std::stoi(argv[i]);
+        } else if (arg == "-ot" || arg == "--override-tensors") {
+            if (++i >= argc) {
+                invalid_param = true;
+                break;
+            }
+            auto value = argv[i];
+            /* static */ std::map<std::string, ggml_backend_buffer_type_t> buft_list;
+            if (buft_list.empty()) {
+                // enumerate all the devices and add their buffer types to the list
+                for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
+                    auto * dev = ggml_backend_dev_get(i);
+                    auto * buft = ggml_backend_dev_buffer_type(dev);
+                    if (buft) {
+                        buft_list[ggml_backend_buft_name(buft)] = buft;
+                    }
+                }
+            }
+            auto override_span_len = std::strcspn(value, ",");
+            while (override_span_len > 0) {
+                // Stamps null terminators into the argv
+                // value for this option to avoid the
+                // memory leak present in the implementation
+                // over in arg.cpp. Maybe allowable because we
+                // only parse these args once in this program.
+                auto override = value;
+                if (value[override_span_len] != '\0') {
+                    value[override_span_len] = '\0';
+                    value = &value[override_span_len + 1];
+                } else {
+                    value = &value[override_span_len];
+                }
+                auto tensor_name_span_len = std::strcspn(override, "=");
+                if (tensor_name_span_len >= override_span_len) {
+                    invalid_param = true;
+                    break;
+                }
+                override[tensor_name_span_len] = '\0';
+                auto tensor_name = override;
+                auto buffer_type = &override[tensor_name_span_len + 1];
+                if (buft_list.find(buffer_type) == buft_list.end()) {
+                    printf("Available buffer types:\n");
+                    for (const auto & it : buft_list) {
+                        printf("  %s\n", ggml_backend_buft_name(it.second));
+                    }
+                    invalid_param = true;
+                    break;
+                }
+                params.tensor_buft_overrides.push_back({tensor_name, buft_list.at(buffer_type)});
+                override_span_len = std::strcspn(value, ",");
+            }
         } else if (arg == "-o" || arg == "--output") {
             if (++i >= argc) {
                 invalid_param = true;
@@ -667,6 +720,11 @@ static cmd_params parse_cmd_params(int argc, char ** argv) {
         params.poll = cmd_params_defaults.poll;
     }
 
+    // Attach terminators to options that requre them
+    if (!params.tensor_buft_overrides.empty()) {
+        params.tensor_buft_overrides.push_back({nullptr, nullptr});
+    }
+
     return params;
 }
 
@@ -689,6 +747,7 @@ struct cmd_params_instance {
     bool               no_kv_offload;
     bool               flash_attn;
     std::vector<float> tensor_split;
+    std::vector<llama_model_tensor_buft_override> tensor_buft_overrides;
     bool               use_mmap;
     bool               embeddings;
 
@@ -733,6 +792,13 @@ struct cmd_params_instance {
         mparams.tensor_split = tensor_split.data();
         mparams.use_mmap     = use_mmap;
 
+        if (tensor_buft_overrides.empty()) {
+            mparams.tensor_buft_overrides = nullptr;
+        } else {
+            GGML_ASSERT(tensor_buft_overrides.back().pattern == nullptr && "Tensor buffer overrides not terminated with empty pattern");
+            mparams.tensor_buft_overrides = tensor_buft_overrides.data();
+        }
+
         return mparams;
     }
 
@@ -804,6 +870,7 @@ static std::vector<cmd_params_instance> get_cmd_params_instances(const cmd_param
                 /* .no_kv_offload= */ nkvo,
                 /* .flash_attn   = */ fa,
                 /* .tensor_split = */ ts,
+                /* .tensor_buft_overrides = */ params.tensor_buft_overrides,
                 /* .use_mmap     = */ mmp,
                 /* .embeddings   = */ embd,
             };
@@ -833,6 +900,7 @@ static std::vector<cmd_params_instance> get_cmd_params_instances(const cmd_param
                 /* .no_kv_offload= */ nkvo,
                 /* .flash_attn   = */ fa,
                 /* .tensor_split = */ ts,
+                /* .tensor_buft_overrides = */ params.tensor_buft_overrides,
                 /* .use_mmap     = */ mmp,
                 /* .embeddings   = */ embd,
             };
@@ -862,6 +930,7 @@ static std::vector<cmd_params_instance> get_cmd_params_instances(const cmd_param
                 /* .no_kv_offload= */ nkvo,
                 /* .flash_attn   = */ fa,
                 /* .tensor_split = */ ts,
+                /* .tensor_buft_overrides = */ params.tensor_buft_overrides,
                 /* .use_mmap     = */ mmp,
                 /* .embeddings   = */ embd,
             };