vulkan: implement ROUND

Signed-off-by: Giuseppe Scrivano <gscrivan@redhat.com>

Author: giuseppe

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

@@ -666,6 +666,7 @@ struct vk_device_struct {
     vk_pipeline pipeline_abs[2];
     vk_pipeline pipeline_softplus[2];
     vk_pipeline pipeline_step[2];
+    vk_pipeline pipeline_round[2];
 
     vk_pipeline pipeline_add1_f16_f16;
     vk_pipeline pipeline_add1_f16_f32;
@@ -3832,6 +3833,7 @@ static void ggml_vk_load_shaders(vk_device& device) {
     CREATE_UNARY(abs)
     CREATE_UNARY(softplus)
     CREATE_UNARY(step)
+    CREATE_UNARY(round)
 #undef CREATE_UNARY
 
 #define CREATE_UNARY_RTE(name)  \
@@ -8260,6 +8262,8 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
                 return ctx->device->pipeline_softplus[dst->type == GGML_TYPE_F16];
             case GGML_UNARY_OP_STEP:
                 return ctx->device->pipeline_step[dst->type == GGML_TYPE_F16];
+            case GGML_UNARY_OP_ROUND:
+                return ctx->device->pipeline_round[dst->type == GGML_TYPE_F16];
             default:
                 break;
         }
@@ -11283,6 +11287,7 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
         case GGML_UNARY_OP_ABS:
         case GGML_UNARY_OP_SOFTPLUS:
         case GGML_UNARY_OP_STEP:
+        case GGML_UNARY_OP_ROUND:
             break;
         default:
             return false;
@@ -11637,6 +11642,7 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
         case GGML_UNARY_OP_ABS:
         case GGML_UNARY_OP_SOFTPLUS:
         case GGML_UNARY_OP_STEP:
+        case GGML_UNARY_OP_ROUND:
             ggml_vk_unary(ctx, compute_ctx, src0, node);
             break;
         default:
@@ -11915,6 +11921,7 @@ static bool ggml_vk_compute_forward(ggml_backend_vk_context * ctx, ggml_cgraph *
         case GGML_UNARY_OP_ABS:
         case GGML_UNARY_OP_SOFTPLUS:
         case GGML_UNARY_OP_STEP:
+        case GGML_UNARY_OP_ROUND:
             buf = tensor->buffer;
             break;
         default:
@@ -13519,6 +13526,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
                 case GGML_UNARY_OP_ABS:
                 case GGML_UNARY_OP_SOFTPLUS:
                 case GGML_UNARY_OP_STEP:
+                case GGML_UNARY_OP_ROUND:
                     return ggml_is_contiguous(op->src[0]) &&
                            (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16) &&
                            (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) &&
@@ -14437,6 +14445,9 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
             case GGML_UNARY_OP_STEP:
                 tensor_clone = ggml_step(ggml_ctx, src_clone[0]);
                 break;
+            case GGML_UNARY_OP_ROUND:
+                tensor_clone = ggml_round(ggml_ctx, src_clone[0]);
+                break;
             default:
                 std::cerr << "Missing vk_check_results OP: " << ggml_op_name(tensor->op) << std::endl;
                 GGML_ABORT("fatal error");

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

@@ -0,0 +1,29 @@
+#version 450
+
+#include "generic_head.glsl"
+#include "types.glsl"
+
+#extension GL_EXT_control_flow_attributes : enable
+
+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer X {A_TYPE data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];};
+
+void main() {
+    const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x;
+
+    if (i >= p.KX) {
+        return;
+    }
+
+    const float x = float(data_a[i]);
+    float result;
+    // Round halfway cases away from zero as roundf does.
+    if (x >= 0.0) {
+        result = floor(x + 0.5);
+    } else {
+        result = ceil(x - 0.5);
+    }
+    data_d[i] = D_TYPE(result);
+}

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

@@ -853,6 +853,8 @@ void process_shaders() {
     string_to_spv("fill_f32",       "fill.comp",        {{"D_TYPE", "float"},       {"FLOAT_TYPE", "float"}});
     string_to_spv("step_f16",       "step.comp",        {{"A_TYPE", "float16_t"},   {"D_TYPE", "float16_t"}});
     string_to_spv("step_f32",       "step.comp",        {{"A_TYPE", "float"},       {"D_TYPE", "float"}});
+    string_to_spv("round_f16",      "round.comp",       {{"A_TYPE", "float16_t"},   {"D_TYPE", "float16_t"}});
+    string_to_spv("round_f32",      "round.comp",       {{"A_TYPE", "float"},       {"D_TYPE", "float"}});
 
     for (auto rte : {false, true}) {
         std::string suffix = rte ? "_rte" : "";