Reference "Software Emulation" implementation of GGML_OP_MUL_MAT

[wilderfield]

Today I was wanting a simple reference implementation of ggml_compute_forward_mul_mat().

Not a hyper optimized BLAS/AVX version.

Purpose == sanity checking, making sure I understand the compute exactly.

I hoped I would find it in the tests, but I haven't stumbled upon it yet.

Does one exist?

I tried to write my own, but pretty sure its incorrect, as I was seeing strange behavior at runtime, and I expect memory corruption.

Looking at the shapes I'm seeing, it appears that the weights buffers are already transposed.

Example shapes:
float *x; // Input with shape (4096, 2, 1, 1)
float *y; // Weights with shape (4096, 11008, 1, 1)
float *z; // Output with shape (11008, 2, 1, 1)

Note how the matching inner dimension appears to be ne10.

void ggml_ref_mul_mat(const struct ggml_tensor * src0, const struct ggml_tensor * src1, struct ggml_tensor * dst, void * wdata, size_t wsize) {
    GGML_ASSERT(ggml_ref_can_mul_mat(src0, src1, dst));
    LOG("ggml_compute_forward_mul_mat called\n");
    LOG("backend type: %d\n", dst->backend);
    LOG("src0: %s dtype: %d (%d,%d,%d,%d)\n", src0->name, src0->type, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3]);
    LOG("src1: %s dtype: %d (%d,%d,%d,%d)\n", src1->name, src1->type, src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3]);
    LOG("dst: %s dtype: %d (%d,%d,%d,%d)\n", dst->name, dst->type, dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);

    GGML_TENSOR_BINARY_OP_LOCALS

    ggml_to_float_t  const to_float = (ggml_to_float_t) dequantize_row_q4_K;

    void* w = src0->data;

    std::vector<float> dequantized_weights(ne00*ne01*ne02*ne03);
    float* y = dequantized_weights.data();

    // Attempt to convert Q4_K to floats
    for (int64_t i03 = 0; i03 < ne03; ++i03) {
        for (int64_t i02 = 0; i02 < ne02; ++i02) {
            for (int64_t i01 = 0; i01 < ne01; ++i01) {
                to_float((const char *) w + i01*nb01 + i02*nb02 + i03*nb03, y + i01*ne00 + i02*ne01 + i03*ne02, ne00);
            }
        }
    }

    float* x = (float*)(src1->data);

    float* z = (float*)(dst->data);

    ggml_vec_set_f32(ne0*ne1*ne2*ne3, z, 0);

    for (int i = 0; i < ne3; ++i) {
        for (int j = 0; j < ne2; ++j) {
            for (int k = 0; k < ne1; ++k) {
                for (int l = 0; l < ne0; ++l) {
                    for (int m = 0; m < ne00; ++m) { // Shared inner dimension
                        int indexZ = (((i * ne2 + j) * ne1 + k) * ne0) + l;
                        int indexX = (((i * ne2 + j) * ne1 + k) * ne00) + m;
                        int indexY = (((i * ne2 + j) * ne0 + l) * ne00) + m;
                        z[indexZ] += x[indexX] * y[indexY];
                        // z[i][j][k][l] += x[i][j][k][m] * y[i][j][l][m];
                    }
                }
            }
        }
    }

}

[ggerganov]

Take a look at this example: ggml-org/ggml#713

[wilderfield]

Thanks @ggerganov. I found a bug in that example, and filed a PR: ggml-org/ggml#770

However, it worked as the perfect testbench for me to fool around until I understood something.

GGML mul_mat computes:

$$ A * B^T = C^T $$

$$ (m x k) * (n x k) = (n x m) $$

Here is my functioning emulation code:

    // Compute C^T
    // The input matricies A and B have the same width for cache performance
    // For some reason, the result of the computation is C^T
    for (int i3 = 0; i3 < ne3; ++i3) {
        for (int i2 = 0; i2 < ne2; ++i2) {
            for (int i1 = 0; i1 < ne1; ++i1) { // rows of C, rows of B
                for (int i0 = 0; i0 < ne0; ++i0) { // cols of C, rows of A
                    cptr[i3*ne2*ne1*ne0 + i2*ne1*ne0 + i1*ne0 + i0] = 0;
                    for (int k = 0; k < ne10; ++k) { // Shared dimension
                        cptr[i3*ne2*ne1*ne0 + i2*ne1*ne0 + i1*ne0 + i0] += \
                        aptr[i3*ne02*ne01*ne00 + i2*ne01*ne00 + i0*ne00 + k] * \
                        bptr[i3*ne12*ne11*ne10 + i2*ne11*ne10 + i1*ne10 + k];
                    }
                }
            }
        }
    }