Function ggml_compute_forward_flash_ff_f16

ggml.c:12585–12716

// ggml_compute_forward_flash_ff

static void ggml_compute_forward_flash_ff_f16(
        const struct ggml_compute_params * params,
        const struct ggml_tensor * a,  // F16
        const struct ggml_tensor * b0, // F16 fc_w
        const struct ggml_tensor * b1, // F32 fc_b
        const struct ggml_tensor * c0, // F16 proj_w
        const struct ggml_tensor * c1, // F32 proj_b
        struct ggml_tensor * dst) {
    int64_t t0 = ggml_perf_time_us();
    UNUSED(t0);

    GGML_TENSOR_LOCALS(int64_t, nea,  a,   ne)
    GGML_TENSOR_LOCALS(size_t,  nba,  a,   nb)
    GGML_TENSOR_LOCALS(int64_t, neb0, b0,  ne)
    GGML_TENSOR_LOCALS(size_t,  nbb0, b0,  nb)
    GGML_TENSOR_LOCALS(int64_t, neb1, b1,  ne)
    GGML_TENSOR_LOCALS(size_t,  nbb1, b1,  nb)
    GGML_TENSOR_LOCALS(int64_t, nec0, c0,  ne)
    GGML_TENSOR_LOCALS(size_t,  nbc0, c0,  nb)
    GGML_TENSOR_LOCALS(int64_t, nec1, c1,  ne)
    GGML_TENSOR_LOCALS(size_t,  nbc1, c1,  nb)
    GGML_TENSOR_LOCALS(int64_t, ne,   dst, ne)
    GGML_TENSOR_LOCALS(size_t,  nb,   dst, nb)

    const int ith = params->ith;
    const int nth = params->nth;

    const int64_t D = nea0;
    //const int64_t N = nea1;
    const int64_t M = neb01;

    GGML_ASSERT(ne0 == nea0);
    GGML_ASSERT(ne1 == nea1);
    GGML_ASSERT(ne2 == nea2);

    GGML_ASSERT(nba0  == sizeof(ggml_fp16_t));
    GGML_ASSERT(nbb00 == sizeof(ggml_fp16_t));
    GGML_ASSERT(nbb10 == sizeof(float));
    GGML_ASSERT(nbc00 == sizeof(ggml_fp16_t));
    GGML_ASSERT(nbc10 == sizeof(float));

    GGML_ASSERT(neb00 == D);
    GGML_ASSERT(neb01 == M);
    GGML_ASSERT(neb10 == M);
    GGML_ASSERT(neb11 == 1);

    GGML_ASSERT(nec00 == M);
    GGML_ASSERT(nec01 == D);
    GGML_ASSERT(nec10 == D);
    GGML_ASSERT(nec11 == 1);

    // dst cannot be transposed or permuted
    GGML_ASSERT(nb0 == sizeof(float));
    GGML_ASSERT(nb0 <= nb1);
    GGML_ASSERT(nb1 <= nb2);
    GGML_ASSERT(nb2 <= nb3);

    if (params->type == GGML_TASK_INIT) {