Writing Vulkan Shaders

#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 G {A_TYPE data_g[];};
layout (binding = 1) readonly buffer X {B_TYPE data_x[];};
layout (binding = 2) writeonly buffer D {D_TYPE data_d[];};

void main() {
    const float GELU_COEF_A    = 0.044715f;
    const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f;

    const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x;

    if (i >= p.KX) {
        return;
    }

    const float xi = float(data_x[i]);
    const float u  = SQRT_2_OVER_PI*xi*(1.0f + GELU_COEF_A*xi*xi);
    const float th = 1.0f - 2.0f / (exp(2.0f*u) + 1.0f);

    const float dg = 0.5f*(1.0f + th) +
                     0.5f*xi*(1.0f - th*th)*SQRT_2_OVER_PI*(1.0f + 3.0f*GELU_COEF_A*xi*xi);

    data_d[i] = D_TYPE(float(data_g[i]) * dg);
}

string_to_spv("gelu_back_f32",  "gelu_back.comp",   {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}});

vk_pipeline pipeline_gelu_back_f32;

ggml_vk_create_pipeline(device, device->pipeline_gelu_back_f32, "gelu_back_f32",
                        gelu_back_f32_len, gelu_back_f32_data, "main",
                        3,                              // number of bindings
                        sizeof(vk_op_push_constants),   // push-constant size
                        {512, 1, 1},                    // workgroup denominators
                        {}, 1);

case GGML_OP_GELU_BACK:
    if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
        return ctx->device->pipeline_gelu_back_f32;
    }
    return nullptr;

static void ggml_vk_gelu_back(ggml_backend_vk_context * ctx, vk_context& subctx,
                              const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
    ggml_vk_op_f32<vk_op_push_constants>(ctx, subctx, src0, src1, nullptr, nullptr, dst,
        GGML_OP_GELU_BACK, { (uint32_t)ggml_nelements(src0), 0, 0.0f, 0.0f, 0.0f, 0.0f });
}

case GGML_OP_GELU_BACK:
    ggml_vk_gelu_back(ctx, compute_ctx, src0, src1, node);
    break;

case GGML_OP_GELU_BACK:
    return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;

library(ggmlR)

run_gelu_back <- function(x, dy, backend = c("cpu", "vulkan")) {
  backend <- match.arg(backend)

  ctx <- ggml_init(16 * 1024 * 1024)
  ggml_set_no_alloc(ctx, TRUE)
  on.exit(ggml_free(ctx), add = TRUE)

  if (backend == "vulkan") {
    b <- ggml_vulkan_init(0L)
  } else {
    b <- ggml_backend_cpu_init()
    ggml_backend_cpu_set_n_threads(b, 2L)
  }
  on.exit(ggml_backend_free(b), add = TRUE)

  g <- ggml_new_tensor_1d(ctx, GGML_TYPE_F32, length(dy))
  x_t <- ggml_new_tensor_1d(ctx, GGML_TYPE_F32, length(x))
  ggml_set_input(g)
  ggml_set_input(x_t)

  out <- ggml_gelu_back(ctx, g, x_t)
  ggml_set_output(out)

  # Does this device run the node itself, or would it fall back to the CPU?
  on_device <- ggml_backend_dev_supports_op(ggml_backend_get_device(b), out)

  buf <- ggml_backend_alloc_ctx_tensors(ctx, b)
  ggml_backend_tensor_set_data(g, dy)
  ggml_backend_tensor_set_data(x_t, x)
  ggml_backend_graph_compute(b, ggml_build_forward_expand(ctx, out))

  list(value = ggml_backend_tensor_get_data(out), on_device = on_device)
}

set.seed(1)
x  <- rnorm(10000, sd = 3)
dy <- rnorm(10000)

cpu <- run_gelu_back(x, dy, "cpu")
cpu$on_device

if (ggml_vulkan_available()) {
  gpu <- run_gelu_back(x, dy, "vulkan")
  print(gpu$on_device)                # TRUE: the shader runs, not a fallback
  max(abs(gpu$value - cpu$value))     # float32 rounding only
}