Repository: ggml-org/llama.cpp
Pull request: vulkan: add exp operation

*Make sure to read the [contributing guidelines](https://github.com/ggml-org/llama.cpp/blob/master/CONTRIBUTING.md) before submitting a PR*

We have found that exp operation is not supported while implementing Vocos, a neural vocoder. The exp operation implementation in Vulkan results in expected output from our model, and the test results on Apple M4 MAX and NVIDIA A6000 are as follows.

test-backend-ops results on Apple M4 Max
```
./build/bin/test-backend-ops -o EXP
ggml_vulkan: Found 1 Vulkan devices:
ggml_vulkan: 0 = Apple M4 Max (MoltenVK) | uma: 1 | fp16: 1 | bf16: 0 | warp size: 32 | shared memory: 32768 | int dot: 0 | matrix cores: none
Testing 4 devices
...
Backend 1/4: Metal
  Device description: Apple M4 Max
  Device memory: 27648 MB (27642 MB free)

  EXP(type=f16,ne_a=[128,2,2,2],v=0): not supported [Metal]
  EXP(type=f16,ne_a=[5,7,11,13],v=0): not supported [Metal]
  EXP(type=f16,ne_a=[128,2,2,2],v=1): not supported [Metal]
  EXP(type=f16,ne_a=[5,7,11,13],v=1): not supported [Metal]
  EXP(type=f32,ne_a=[128,2,2,2],v=0): OK
  EXP(type=f32,ne_a=[5,7,11,13],v=0): OK
  EXP(type=f32,ne_a=[128,2,2,2],v=1): not supported [Metal]
  EXP(type=f32,ne_a=[5,7,11,13],v=1): not supported [Metal]
  10844/10844 tests passed
  Backend Metal: OK
ggml_metal_free: deallocating
ggml_metal_mem_pool_free: freeing memory pool, num heaps = 0 (total = 0)
ggml_metal_mem_pool_free: freeing memory pool, num heaps = 0 (total = 0)
ggml_metal_mem_pool_free: freeing memory pool, num heaps = 0 (total = 0)
ggml_metal_mem_pool_free: freeing memory pool, num heaps = 0 (total = 0)
ggml_metal_mem_pool_free: freeing memory pool, num heaps = 0 (total = 0)
ggml_metal_mem_pool_free: freeing memory pool, num heaps = 0 (total = 0)
ggml_metal_mem_pool_free: freeing memory pool, num heaps = 0 (total = 0)
ggml_metal_mem_pool_free: freeing memory pool, num heaps = 0 (total = 0)
Backend 2/4: Vulkan0
  Device description: Apple M4 Max
  Device memory: 36864 MB (36864 MB free)

  EXP(type=f16,ne_a=[128,2,2,2],v=0): OK
  EXP(type=f16,ne_a=[5,7,11,13],v=0): OK
  EXP(type=f16,ne_a=[128,2,2,2],v=1): not supported [Vulkan0]
  EXP(type=f16,ne_a=[5,7,11,13],v=1): not supported [Vulkan0]
  EXP(type=f32,ne_a=[128,2,2,2],v=0): OK
  EXP(type=f32,ne_a=[5,7,11,13],v=0): OK
  EXP(type=f32,ne_a=[128,2,2,2],v=1): not supported [Vulkan0]
  EXP(type=f32,ne_a=[5,7,11,13],v=1): not supported [Vulkan0]
  10844/10844 tests passed
  Backend Vulkan0: OK
Backend 3/4: BLAS
  Device description: Accelerate
  Device memory: 0 MB (0 MB free)

  EXP(type=f16,ne_a=[128,2,2,2],v=0): not supported [BLAS]
  EXP(type=f16,ne_a=[5,7,11,13],v=0): not supported [BLAS]
  EXP(type=f16,ne_a=[128,2,2,2],v=1): not supported [BLAS]
  EXP(type=f16,ne_a=[5,7,11,13],v=1): not supported [BLAS]
  EXP(type=f32,ne_a=[128,2,2,2],v=0): not supported [BLAS]
  EXP(type=f32,ne_a=[5,7,11,13],v=0): not supported [BLAS]
  EXP(type=f32,ne_a=[128,2,2,2],v=1): not supported [BLAS]
  EXP(type=f32,ne_a=[5,7,11,13],v=1): not supported [BLAS]
  10844/10844 tests passed
  Backend BLAS: OK
Backend 4/4: CPU
  Skipping CPU backend
4/4 backends passed
OK
```

test-backend-ops results on NVIDIA RTX A6000
```
./build/bin/test-backend-ops -o EXP
ggml_cuda_init: GGML_CUDA_FORCE_MMQ:    no
ggml_cuda_init: GGML_CUDA_FORCE_CUBLAS: no
ggml_cuda_init: found 1 CUDA devices:
  Device 0: NVIDIA RTX A6000, compute capability 8.6, VMM: yes
register_backend: registered backend CUDA (1 devices)
register_device: registered device CUDA0 (NVIDIA RTX A6000)
ggml_vulkan: Found 1 Vulkan devices:
ggml_vulkan: 0 = NVIDIA RTX A6000 (NVIDIA) | uma: 0 | fp16: 1 | bf16: 0 | warp size: 32 | shared memory: 49152 | int dot: 1 | matrix cores: NV_coopmat2
register_backend: registered backend Vulkan (1 devices)
register_device: registered device Vulkan0 (NVIDIA RTX A6000)
register_backend: registered backend CPU (1 devices)
register_device: registered device CPU (AMD Ryzen Threadripper PRO 5975WX 32-Cores)
load_backend: failed to find ggml_backend_init in /home/dongwon/prj/llama.cpp/build/bin/libggml-cuda.so
load_backend: failed to find ggml_backend_init in /home/dongwon/prj/llama.cpp/build/bin/libggml-vulkan.so
load_backend: failed to find ggml_backend_init in /home/dongwon/prj/llama.cpp/build/bin/libggml-cpu.so
Testing 3 devices

Backend 1/3: CUDA0
  Device description: NVIDIA RTX A6000
  Device memory: 48539 MB (48223 MB free)

update_cuda_graph_executable: CUDA graph update failed
ggml_backend_cuda_graph_compute: disabling CUDA graphs due to too many consecutive updates
  EXP(type=f16,ne_a=[128,2,2,2],v=0): OK
  EXP(type=f16,ne_a=[5,7,11,13],v=0): OK
  EXP(type=f16,ne_a=[128,2,2,2],v=1): not supported [CUDA0] 
  EXP(type=f16,ne_a=[5,7,11,13],v=1): not supported [CUDA0] 
  EXP(type=f32,ne_a=[128,2,2,2],v=0): OK
  EXP(type=f32,ne_a=[5,7,11,13],v=0): OK
  EXP(type=f32,ne_a=[128,2,2,2],v=1): not supported [CUDA0] 
  EXP(type=f32,ne_a=[5,7,11,13],v=1): not supported [CUDA0] 
  10844/10844 tests passed
  Backend CUDA0: OK
Backend 2/3: Vulkan0
  Device description: NVIDIA RTX A6000
  Device memory: 49140 MB (49140 MB free)

  EXP(type=f16,ne_a=[128,2,2,2],v=0): OK
  EXP(type=f16,ne_a=[5,7,11,13],v=0): OK
  EXP(type=f16,ne_a=[128,2,2,2],v=1): not supported [Vulkan0] 
  EXP(type=f16,ne_a=[5,7,11,13],v=1): not supported [Vulkan0] 
  EXP(type=f32,ne_a=[128,2,2,2],v=0): OK
  EXP(type=f32,ne_a=[5,7,11,13],v=0): OK
  EXP(type=f32,ne_a=[128,2,2,2],v=1): not supported [Vulkan0] 
  EXP(type=f32,ne_a=[5,7,11,13],v=1): not supported [Vulkan0] 
  10844/10844 tests passed
  Backend Vulkan0: OK
Backend 3/3: CPU
  Skipping CPU backend
3/3 backends passed
OK
```

Changed files:
filename: ggml/src/ggml-vulkan/ggml-vulkan.cpp
patch: @@ -490,6 +490,7 @@ struct vk_device_struct {
     vk_pipeline pipeline_l2_norm_f32;
 
     // [src/dst 0=fp32,1=fp16]
+    vk_pipeline pipeline_exp[2];
     vk_pipeline pipeline_gelu[2];
     vk_pipeline pipeline_gelu_erf[2];
     vk_pipeline pipeline_gelu_quick[2];
@@ -3066,6 +3067,7 @@ static void ggml_vk_load_shaders(vk_device& device) {
     ggml_vk_create_pipeline(device, device->pipeline_ ## name [0], #name "_f32", name ## _f32_len, name ## _f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1);  \
     ggml_vk_create_pipeline(device, device->pipeline_ ## name [1], #name "_f16", name ## _f16_len, name ## _f16_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1);
 
+    CREATE_UNARY(exp)
     CREATE_UNARY(gelu)
     CREATE_UNARY(gelu_erf)
     CREATE_UNARY(gelu_quick)
@@ -7133,6 +7135,8 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
         }
 
         switch (ggml_get_unary_op(dst)) {
+            case GGML_UNARY_OP_EXP:
+                return ctx->device->pipeline_exp[dst->type == GGML_TYPE_F16];
             case GGML_UNARY_OP_SILU:
                 return ctx->device->pipeline_silu[dst->type == GGML_TYPE_F16];
             case GGML_UNARY_OP_GELU:
@@ -9738,6 +9742,7 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
         return false;
     case GGML_OP_UNARY:
         switch (ggml_get_unary_op(node)) {
+        case GGML_UNARY_OP_EXP:
         case GGML_UNARY_OP_SILU:
         case GGML_UNARY_OP_GELU:
         case GGML_UNARY_OP_GELU_ERF:
@@ -10015,6 +10020,7 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
         break;
     case GGML_OP_UNARY:
         switch (ggml_get_unary_op(node)) {
+        case GGML_UNARY_OP_EXP:
         case GGML_UNARY_OP_SILU:
         case GGML_UNARY_OP_GELU:
         case GGML_UNARY_OP_GELU_ERF:
@@ -10251,6 +10257,7 @@ static bool ggml_vk_compute_forward(ggml_backend_vk_context * ctx, ggml_cgraph *
         break;
     case GGML_OP_UNARY:
         switch (ggml_get_unary_op(tensor)) {
+        case GGML_UNARY_OP_EXP:
         case GGML_UNARY_OP_SILU:
         case GGML_UNARY_OP_GELU:
         case GGML_UNARY_OP_GELU_ERF:
@@ -11166,6 +11173,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
     switch (op->op) {
         case GGML_OP_UNARY:
             switch (ggml_get_unary_op(op)) {
+                case GGML_UNARY_OP_EXP:
                 case GGML_UNARY_OP_GELU:
                 case GGML_UNARY_OP_GELU_ERF:
                 case GGML_UNARY_OP_GELU_QUICK:
@@ -11965,6 +11973,9 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
         }
     } else if (tensor->op == GGML_OP_UNARY) {
         switch (ggml_get_unary_op(tensor)) {
+        case GGML_UNARY_OP_EXP:
+            tensor_clone = ggml_exp(ggml_ctx, src_clone[0]);
+            break;
         case GGML_UNARY_OP_SILU:
             tensor_clone = ggml_silu(ggml_ctx, src_clone[0]);
             break;

filename: ggml/src/ggml-vulkan/vulkan-shaders/exp.comp
patch: @@ -0,0 +1,20 @@
+#version 450
+
+#include "generic_head.comp"
+#include "types.comp"
+
+#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;
+    }
+    data_d[i] = D_TYPE(exp(float(data_a[i])));
+}

filename: ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
patch: @@ -586,6 +586,8 @@ void process_shaders() {
 
     string_to_spv("upscale_f32", "upscale.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}});
 
+    string_to_spv("exp_f16",        "exp.comp",         {{"A_TYPE", "float16_t"},   {"D_TYPE", "float16_t"}});
+    string_to_spv("exp_f32",        "exp.comp",         {{"A_TYPE", "float"},       {"D_TYPE", "float"}});
     string_to_spv("gelu_f16",       "gelu.comp",        {{"A_TYPE", "float16_t"},   {"D_TYPE", "float16_t"}});
     string_to_spv("gelu_f32",       "gelu.comp",        {{"A_TYPE", "float"},       {"D_TYPE", "float"}});
     string_to_spv("gelu_erf_f16",   "gelu_erf.comp",    {{"A_TYPE", "float16_t"},   {"D_TYPE", "float16_t"}});