[DRAFT]: Formalizing scheduler and async backend behavior through tests

ggml/include/ggml-rpc.h

@@ -11,7 +11,7 @@ extern "C" {
 #define RPC_PROTO_PATCH_VERSION    0
 
 #ifdef  __cplusplus
-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
+static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
 #endif
 
 #define GGML_RPC_MAX_SERVERS       16

ggml/include/ggml.h

@@ -590,6 +590,8 @@ extern "C" {
 
         GGML_OP_GLU,
 
+        GGML_OP_SLEEP,
+
         GGML_OP_COUNT,
     };
 
@@ -2722,6 +2724,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

@@ -1997,6 +1997,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);
@@ -2265,6 +2269,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
@@ -2288,6 +2289,27 @@ 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) {
+    const ggml_tensor * src0 = dst->src[0];
+
+    if (params->ith != 0) {
+        return;
+    }
+
+    GGML_ASSERT(src0->type == dst->type);
+    GGML_ASSERT(ggml_are_same_shape(src0, dst));
+    GGML_ASSERT(ggml_is_contiguous(src0));
+    GGML_ASSERT(ggml_is_contiguous(dst));
+
+    const auto t_end = std::chrono::steady_clock::now() + std::chrono::microseconds(ggml_get_op_params_i32(dst, 0));
+
+    memcpy(dst->data, src0->data, ggml_nbytes(dst));
+
+    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

@@ -87,6 +87,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

@@ -42,6 +42,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"
@@ -2392,6 +2393,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;
@@ -5297,6 +5301,12 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
             return true;
         case GGML_OP_LIGHTNING_INDEXER:
             return ggml_cuda_lightning_indexer_supported(dev_ctx->device, op);
+        case GGML_OP_SLEEP:
+#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
+            return false;
+#else
+            return op->type == op->src[0]->type && ggml_is_contiguous(op->src[0]);
+#endif
 
         default:
             return false;

ggml/src/ggml-cuda/sleep.cu

@@ -0,0 +1,41 @@
+#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 following memcpy on the same stream cannot start before this kernel 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) {
+    const ggml_tensor * src0 = dst->src[0];
+
+    GGML_ASSERT(src0->type == dst->type);
+    GGML_ASSERT(ggml_are_same_shape(src0, dst));
+    GGML_ASSERT(ggml_is_contiguous(src0));
+    GGML_ASSERT(ggml_is_contiguous(dst));
+
+    cudaStream_t stream = ctx.stream();
+
+    sleep_ns<<<1, 1, 0, stream>>>(1000*(uint64_t) ggml_get_op_params_i32(dst, 0));
+    CUDA_CHECK(cudaMemcpyAsync(dst->data, src0->data, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream));
+}
+
+#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

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/CMakeLists.txt

@@ -109,6 +109,12 @@ if (Vulkan_FOUND)
         "GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT"
     )
 
+    test_shader_extension_support(
+        "GL_EXT_shader_realtime_clock"
+        "${CMAKE_CURRENT_SOURCE_DIR}/vulkan-shaders/feature-tests/shader_clock.comp"
+        "GGML_VULKAN_SHADER_CLOCK_GLSLC_SUPPORT"
+    )
+
     target_link_libraries(ggml-vulkan PRIVATE Vulkan::Vulkan)
     target_include_directories(ggml-vulkan PRIVATE ${CMAKE_CURRENT_BINARY_DIR})
 

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

@@ -818,6 +818,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;
@@ -1013,6 +1014,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];
@@ -1394,6 +1397,11 @@ struct vk_op_push_constants {
     float param4;
 };
 
+struct vk_op_sleep_push_constants {
+    uint32_t ne;
+    uint32_t ticks;
+};
+
 struct vk_op_fwht_push_constants {
     uint32_t n_rows;
     uint32_t src_offset;
@@ -5675,6 +5683,12 @@ static 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 defined(GGML_VULKAN_SHADER_CLOCK_GLSLC_SUPPORT)
+    if (device->shader_clock) {
+        ggml_vk_create_pipeline(device, device->pipeline_sleep, "sleep", sleep_len, sleep_data, "main", 2, sizeof(vk_op_sleep_push_constants), {512, 1, 1}, {}, 1);
+    }
+#endif
+
 #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);
@@ -6203,6 +6217,7 @@ static 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) {
@@ -6252,6 +6267,11 @@ static vk_device ggml_vk_get_device(size_t idx) {
 #if defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
             } else if (strcmp(VK_EXT_SHADER_FLOAT8_EXTENSION_NAME, properties.extensionName) == 0) {
                 shader_float8_extension = true;
+#endif
+#if defined(GGML_VULKAN_SHADER_CLOCK_GLSLC_SUPPORT)
+            } else if (strcmp("VK_KHR_shader_clock", properties.extensionName) == 0 &&
+                       !getenv("GGML_VK_DISABLE_SHADER_CLOCK")) {
+                shader_clock_support = true;
 #endif
             } else if (strcmp("VK_VALVE_shader_mixed_float_dot_product", properties.extensionName) == 0 &&
                        !getenv("GGML_VK_DISABLE_DOT2")) {
@@ -6595,6 +6615,14 @@ static 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) {
@@ -6678,6 +6706,7 @@ static 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;
@@ -11779,6 +11808,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;
@@ -12980,6 +13011,32 @@ static void ggml_vk_fill(ggml_backend_vk_context * ctx, vk_context& subctx, ggml
     ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { dst_buf }, pc, elements);
 }
 
+static void ggml_vk_sleep(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, 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, assume it matches the timestamp counter
+    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)(ggml_nbytes(dst) / sizeof(uint32_t)),
+        (uint32_t) ticks,
+    };
+
+    vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, src0, nullptr, nullptr, dst, GGML_OP_SLEEP);
+    GGML_ASSERT(pipeline != nullptr);
+
+    ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
+    vk_subbuffer src_buf = ggml_vk_tensor_subbuffer(ctx, src0);
+    vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst);
+
+    // dispatch a single workgroup, all of its invocations spin concurrently so the delay does not accumulate
+    std::array<uint32_t, 3> elements = { pipeline->wg_denoms[0], 1, 1 };
+
+    ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { src_buf, dst_buf }, pc, elements);
+}
+
 static 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));
 }
@@ -15386,6 +15443,10 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
     case GGML_OP_FILL:
         ggml_vk_fill(ctx, compute_ctx, node);
 
+        break;
+    case GGML_OP_SLEEP:
+        ggml_vk_sleep(ctx, compute_ctx, src0, node);
+
         break;
     case GGML_OP_SCALE:
         ggml_vk_scale(ctx, compute_ctx, src0, node);
@@ -18349,6 +18410,11 @@ 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:
+            // the shader copies 4-byte units, the delay comes from the device clock
+            return device->shader_clock && op->type == op->src[0]->type &&
+                   ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op) &&
+                   (ggml_nbytes(op) % sizeof(uint32_t)) == 0;
         case GGML_OP_SCALE:
             return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
         case GGML_OP_PAD:
@@ -19114,6 +19180,8 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_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/vulkan-shaders/CMakeLists.txt

@@ -31,6 +31,10 @@ if (GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
     add_compile_definitions(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
     message(STATUS "Enabling E4M3 glslc support")
 endif()
+if (GGML_VULKAN_SHADER_CLOCK_GLSLC_SUPPORT)
+    add_compile_definitions(GGML_VULKAN_SHADER_CLOCK_GLSLC_SUPPORT)
+    message(STATUS "Enabling shader clock glslc support")
+endif()
 if (GGML_VULKAN_SHADER_DEBUG_INFO)
     add_compile_definitions(GGML_VULKAN_SHADER_DEBUG_INFO)
     message(STATUS "Enabling shader debug info")

ggml/src/ggml-vulkan/vulkan-shaders/feature-tests/shader_clock.comp

@@ -0,0 +1,7 @@
+#version 460
+
+#extension GL_EXT_shader_realtime_clock : require
+
+void main()
+{
+}

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

@@ -0,0 +1,36 @@
+#version 450
+
+#extension GL_EXT_shader_realtime_clock : require
+
+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly  buffer A {uint data_a[];};
+layout (binding = 1) writeonly buffer D {uint data_d[];};
+
+layout (push_constant) uniform parameter {
+    uint ne;
+    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() {
+    // the dispatch is a single workgroup, so all invocations spin concurrently and the delay does not accumulate
+    const uvec2 t0 = clockRealtime2x32EXT();
+
+    for (uint iter = 0; iter < MAX_ITER; ++iter) {
+        const uvec2 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;
+        }
+    }
+
+    for (uint i = gl_LocalInvocationID.x; i < p.ne; i += gl_WorkGroupSize.x) {
+        data_d[i] = data_a[i];
+    }
+}