Formalizing scheduler and async backend behavior through tests

ggml/include/ggml-rpc.h

@@ -8,10 +8,10 @@ extern "C" {
 
 #define RPC_PROTO_MAJOR_VERSION    8
 #define RPC_PROTO_MINOR_VERSION    0
-#define RPC_PROTO_PATCH_VERSION    1
+#define RPC_PROTO_PATCH_VERSION    2
 
 #ifdef  __cplusplus
-static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
+static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
 #endif
 
 #define GGML_RPC_MAX_SERVERS       16

ggml/include/ggml.h

@@ -602,6 +602,8 @@ extern "C" {
 
         GGML_OP_GLU,
 
+        GGML_OP_SLEEP,
+
         GGML_OP_COUNT,
     };
 
@@ -2846,6 +2848,12 @@ extern "C" {
             int                   n_tasks,
             void                * userdata);
 
+    // Busy-wait operator for scheduling and backend testing.
+    GGML_API struct ggml_tensor * ggml_sleep(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int32_t               us);
+
     // loss function
 
     GGML_API struct ggml_tensor * ggml_cross_entropy_loss(

ggml/src/ggml-cpu/ggml-cpu.c

@@ -2031,6 +2031,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_fill(params, tensor);
             } break;
+        case GGML_OP_SLEEP:
+            {
+                ggml_compute_forward_sleep(params, tensor);
+            } break;
         case GGML_OP_FLASH_ATTN_EXT:
             {
                 ggml_compute_forward_flash_attn_ext(params, tensor);
@@ -2304,6 +2308,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_REPEAT:
         case GGML_OP_REPEAT_BACK:
         case GGML_OP_LEAKY_RELU:
+        case GGML_OP_SLEEP:
             {
                 n_tasks = 1;
             } break;

ggml/src/ggml-cpu/ops.cpp

@@ -10,6 +10,7 @@
 
 #include <algorithm>
 #include <cfloat>
+#include <chrono>
 #include <cmath>
 
 // ggml_compute_forward_dup
@@ -2352,6 +2353,18 @@ void ggml_compute_forward_fill(const ggml_compute_params * params, ggml_tensor *
     }
 }
 
+// ggml_compute_forward_sleep
+
+void ggml_compute_forward_sleep(const ggml_compute_params * params, ggml_tensor * dst) {
+    if (params->ith != 0) {
+        return;
+    }
+
+    const auto t_end = std::chrono::steady_clock::now() + std::chrono::microseconds(ggml_get_op_params_i32(dst, 0));
+
+    while (std::chrono::steady_clock::now() < t_end) {}
+}
+
 // ggml_compute_tri
 
 static void ggml_compute_forward_tri_f32(const ggml_compute_params * params, ggml_tensor * dst) {

ggml/src/ggml-cpu/ops.h

@@ -86,6 +86,7 @@ void ggml_compute_forward_top_k(const struct ggml_compute_params * params, struc
 void ggml_compute_forward_leaky_relu(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_tri(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_fill(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+void ggml_compute_forward_sleep(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_flash_attn_ext(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_flash_attn_back(
         const struct ggml_compute_params * params,

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -45,6 +45,7 @@
 #include "ggml-cuda/rope.cuh"
 #include "ggml-cuda/roll.cuh"
 #include "ggml-cuda/scale.cuh"
+#include "ggml-cuda/sleep.cuh"
 #include "ggml-cuda/snake.cuh"
 #include "ggml-cuda/softcap.cuh"
 #include "ggml-cuda/softmax.cuh"
@@ -2607,6 +2608,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
         case GGML_OP_FILL:
             ggml_cuda_op_fill(ctx, dst);
             break;
+        case GGML_OP_SLEEP:
+            ggml_cuda_op_sleep(ctx, dst);
+            break;
         case GGML_OP_LIGHTNING_INDEXER:
             ggml_cuda_lightning_indexer(ctx, dst);
             break;
@@ -6109,6 +6113,12 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
             return ggml_cuda_lightning_indexer_supported(dev_ctx->device, op);
         case GGML_OP_TOP_K_MASK:
             return op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_F32;
+        case GGML_OP_SLEEP:
+#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
+            return false;
+#else
+            return true;
+#endif // defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
 
         default:
             return false;

ggml/src/ggml-cuda/sleep.cu

@@ -0,0 +1,32 @@
+#include "sleep.cuh"
+
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+
+// %globaltimer is a nanosecond wall clock, unlike clock64() it is unaffected by the SM clock and by frequency scaling
+static __device__ __forceinline__ uint64_t globaltimer_ns() {
+    uint64_t t;
+    asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t));
+    return t;
+}
+
+// a single thread is enough, the next kernel on the same stream cannot start before this one retires
+static __global__ void sleep_ns(const uint64_t ns) {
+    const uint64_t t0 = globaltimer_ns();
+
+    while (globaltimer_ns() - t0 < ns) {}
+}
+
+void ggml_cuda_op_sleep(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    sleep_ns<<<1, 1, 0, ctx.stream()>>>(1000*(uint64_t) ggml_get_op_params_i32(dst, 0));
+    CUDA_CHECK(cudaGetLastError());
+}
+
+#else
+
+void ggml_cuda_op_sleep(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    GGML_UNUSED(ctx);
+    GGML_UNUSED(dst);
+    GGML_ABORT("GGML_OP_SLEEP requires the %%globaltimer register, which is only available on CUDA");
+}
+
+#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)

ggml/src/ggml-cuda/sleep.cuh

@@ -0,0 +1,3 @@
+#include "common.cuh"
+
+void ggml_cuda_op_sleep(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

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

@@ -993,6 +993,8 @@ void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph
         } else if (tensor->op == GGML_OP_FILL) {
             const float value = ggml_get_op_params_f32(tensor, 0);
             tensor_clone = ggml_fill(ggml_ctx, src_clone[0], value);
+        } else if (tensor->op == GGML_OP_SLEEP) {
+            tensor_clone = ggml_sleep(ggml_ctx, src_clone[0], ggml_get_op_params_i32(tensor, 0));
         } else if (tensor->op == GGML_OP_SQR) {
             tensor_clone = ggml_sqr(ggml_ctx, src_clone[0]);
         } else if (tensor->op == GGML_OP_SQRT) {

ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h

@@ -151,6 +151,10 @@ struct vk_op_push_constants {
     float param4;
 };
 
+struct vk_op_sleep_push_constants {
+    uint32_t ticks;
+};
+
 struct vk_op_fwht_push_constants {
     uint32_t n_rows;
     uint32_t src_offset;

ggml/src/ggml-vulkan/ggml-vulkan-types.h

@@ -710,6 +710,7 @@ struct vk_device_struct {
     bool shader_int64;
     bool buffer_device_address;
     bool vulkan_memory_model;
+    bool shader_clock;
 
     bool add_rms_fusion;
     uint32_t partials_binding_alignment;
@@ -903,6 +904,8 @@ struct vk_device_struct {
     vk_pipeline pipeline_fill_f32;
     vk_pipeline pipeline_fill_f16;
 
+    vk_pipeline pipeline_sleep;
+
     vk_pipeline pipeline_geglu[2];
     vk_pipeline pipeline_reglu[2];
     vk_pipeline pipeline_swiglu[2];

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

@@ -3470,6 +3470,10 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
     ggml_vk_create_pipeline(device, device->pipeline_fill_f32, "fill_f32", fill_f32_len, fill_f32_data, "main", 1, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_fill_f16, "fill_f16", fill_f16_len, fill_f16_data, "main", 1, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1);
 
+    if (device->shader_clock) {
+        ggml_vk_create_pipeline(device, device->pipeline_sleep, "sleep", sleep_len, sleep_data, "main", 1, sizeof(vk_op_sleep_push_constants), {1, 1, 1}, {}, 1);
+    }
+
 #define CREATE_GLU(name)  \
     ggml_vk_create_pipeline(device, device->pipeline_ ## name [0], #name "_f32", name ## _f32_len, name ## _f32_data, "main", 3, sizeof(vk_op_glu_push_constants), {512, 1, 1}, {}, 1, true);   \
     ggml_vk_create_pipeline(device, device->pipeline_ ## name [1], #name "_f16", name ## _f16_len, name ## _f16_data, "main", 3, sizeof(vk_op_glu_push_constants), {512, 1, 1}, {}, 1, true);
@@ -4055,6 +4059,7 @@ vk_device ggml_vk_get_device(size_t idx) {
         bool dot2_f16_support = false;
         bool ocp_microscaling_extension = false;
         bool shader_float8_extension = false;
+        bool shader_clock_support = false;
 
         for (const auto& properties : ext_props) {
             if (strcmp("VK_KHR_maintenance4", properties.extensionName) == 0) {
@@ -4105,6 +4110,9 @@ vk_device ggml_vk_get_device(size_t idx) {
             } else if (strcmp(VK_EXT_SHADER_FLOAT8_EXTENSION_NAME, properties.extensionName) == 0) {
                 shader_float8_extension = true;
 #endif
+            } else if (strcmp("VK_KHR_shader_clock", properties.extensionName) == 0 &&
+                       !getenv("GGML_VK_DISABLE_SHADER_CLOCK")) {
+                shader_clock_support = true;
             } else if (strcmp("VK_VALVE_shader_mixed_float_dot_product", properties.extensionName) == 0 &&
                        !getenv("GGML_VK_DISABLE_DOT2")) {
                 dot2_f16_support = true;
@@ -4443,6 +4451,14 @@ vk_device ggml_vk_get_device(size_t idx) {
             device_extensions.push_back("VK_KHR_shader_integer_dot_product");
         }
 
+        VkPhysicalDeviceShaderClockFeaturesKHR shader_clock_features {};
+        shader_clock_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_CLOCK_FEATURES_KHR;
+        if (shader_clock_support) {
+            last_struct->pNext = (VkBaseOutStructure *)&shader_clock_features;
+            last_struct = (VkBaseOutStructure *)&shader_clock_features;
+            device_extensions.push_back("VK_KHR_shader_clock");
+        }
+
         VkPhysicalDeviceShaderMixedFloatDotProductFeaturesVALVE dot2_features {};
         dot2_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_MIXED_FLOAT_DOT_PRODUCT_FEATURES_VALVE;
         if (dot2_f16_support) {
@@ -4526,6 +4542,7 @@ vk_device ggml_vk_get_device(size_t idx) {
         device->shader_int64 = device_features2.features.shaderInt64;
         device->buffer_device_address = vk12_features.bufferDeviceAddress;
         device->vulkan_memory_model = vk12_features.vulkanMemoryModel;
+        device->shader_clock = shader_clock_support && shader_clock_features.shaderDeviceClock;
 
         if (device->subgroup_size_control) {
             device->subgroup_min_size = subgroup_size_control_props.minSubgroupSize;
@@ -9474,6 +9491,8 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
             return ctx->device->pipeline_arange_f32;
         }
         return nullptr;
+    case GGML_OP_SLEEP:
+        return ctx->device->pipeline_sleep;
     case GGML_OP_FILL:
         if (dst->type == GGML_TYPE_F32) {
             return ctx->device->pipeline_fill_f32;
@@ -10689,6 +10708,31 @@ void ggml_vk_fill(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor
     ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { dst_buf }, pc, elements);
 }
 
+static void ggml_vk_sleep(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) {
+    VK_LOG_DEBUG("ggml_vk_sleep(dst=" << dst << ", us=" << ggml_get_op_params_i32(dst, 0) << ")");
+
+    // Vulkan does not specify the period of the shader realtime clock this assumes it matches the timestamp counter.
+    // A driver where this does not hold could sleep too short or too long by that ratio
+    const float    ns_per_tick = ctx->device->properties.limits.timestampPeriod;
+    const uint64_t ns          = 1000 * (uint64_t) ggml_get_op_params_i32(dst, 0);
+    const uint64_t ticks       = std::min<uint64_t>(ns_per_tick > 0.0f ? (uint64_t)(ns / ns_per_tick) : ns, std::numeric_limits<uint32_t>::max());
+
+    vk_op_sleep_push_constants pc = {
+        (uint32_t) ticks,
+    };
+
+    vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, nullptr, nullptr, nullptr, dst, GGML_OP_SLEEP);
+    GGML_ASSERT(pipeline != nullptr);
+
+    ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
+    vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst);
+
+    // a single invocation spins, the buffer is bound only to keep the spin from being optimized out
+    std::array<uint32_t, 3> elements = { 1, 1, 1 };
+
+    ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { dst_buf }, pc, elements);
+}
+
 void ggml_vk_sin(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) {
     ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_SIN, vk_op_unary_push_constants_init(src0, dst));
 }
@@ -12551,6 +12595,10 @@ bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, in
     case GGML_OP_FILL:
         ggml_vk_fill(ctx, compute_ctx, node);
 
+        break;
+    case GGML_OP_SLEEP:
+        ggml_vk_sleep(ctx, compute_ctx, node);
+
         break;
     case GGML_OP_SCALE:
         if (ctx->fused_hc_post_gate) {
@@ -15912,6 +15960,8 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
             return op->type == GGML_TYPE_F32;
         case GGML_OP_FILL:
             return op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16;
+        case GGML_OP_SLEEP:
+            return device->shader_clock;
         case GGML_OP_SCALE:
             return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
         case GGML_OP_PAD:

ggml/src/ggml-vulkan/vulkan-shaders/sleep.comp

@@ -0,0 +1,37 @@
+#version 450
+
+#extension GL_EXT_shader_realtime_clock : require
+
+layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
+
+// never written, it only anchors the spin below so the compiler cannot drop it
+layout (binding = 0) writeonly buffer D {uint data_d[];};
+
+layout (push_constant) uniform parameter {
+    uint ticks;
+} p;
+
+// bounds the spin in case the clock does not advance, turning a hang into a too-short sleep
+const uint MAX_ITER = 1u << 27;
+
+void main() {
+    const uvec2 t0 = clockRealtime2x32EXT();
+
+    uvec2 t = t0;
+
+    for (uint iter = 0; iter < MAX_ITER; ++iter) {
+        t = clockRealtime2x32EXT();
+
+        const uint lo = t.x - t0.x;
+        const uint hi = t.y - t0.y - (t.x < t0.x ? 1u : 0u);
+
+        if (hi != 0u || lo >= p.ticks) {
+            break;
+        }
+    }
+
+    // the store never runs, but the compiler cannot prove it and has to keep the spin
+    if (t.x == 0xFFFFFFFFu && t.y == 0xFFFFFFFFu) {
+        data_d[0] = t.x;
+    }
+}

ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp

@@ -1039,6 +1039,7 @@ void process_shaders() {
     string_to_spv("arange_f32",     "arange.comp",      {{"A_TYPE", "float"},       {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
     string_to_spv("fill_f32",       "fill.comp",        {{"D_TYPE", "float"},       {"FLOAT_TYPE", "float"}});
     string_to_spv("fill_f16",       "fill.comp",        {{"D_TYPE", "float16_t"},   {"FLOAT_TYPE", "float"}});
+    string_to_spv("sleep",          "sleep.comp",       {});
     string_to_spv("step_f16",       "unary.comp",       {{"A_TYPE", "float16_t"},   {"D_TYPE", "float16_t"}, {"OP", "op_step"}});
     string_to_spv("step_f32",       "unary.comp",       {{"A_TYPE", "float"},       {"D_TYPE", "float"},     {"OP", "op_step"}});
     string_to_spv("round_f16",      "unary.comp",       {{"A_TYPE", "float16_t"},   {"D_TYPE", "float16_t"}, {"OP", "op_round"}});