llama.cppのレジスタ使用について

実現したいこと

レジスタの使用についてまとめたい

前提

https://salykova.github.io/matmul
このサイトにあるKernelの話を基にllama.cppの行列演算部分のtilingサイズを変更していました。
今回使用したマシンではレジスタ数が32個で1つのレジスタに4個のfloatが格納でき、4*29のサイズが最も速度向上を実現できました。
そこで、この条件では32個中31個のレジスタを使用している計算になるため、残り1つのレジスタはどこで使用されているのかをご意見いただきたいです。

現状考えているのは、vec_dotにおける一時的な計算用バッファとして1つ余らせているのかと考えています。
他のマシンでも1つレジスタが余る状態が性能向上しているので、何かしらの理由があるのかと考えています。

発生している問題・エラーメッセージ

エラーメッセージ

該当のソースコード

static void ggml_compute_forward_mul_mat_one_chunk(
const struct ggml_compute_params * params,
struct ggml_tensor * dst,
const enum ggml_type type,
const int64_t num_rows_per_vec_dot,
const int64_t ir0_start,
const int64_t ir0_end,
const int64_t ir1_start,
const int64_t ir1_end,
const int mr,
const int nr) {

const struct ggml_tensor * src0 = dst->src[0];
const struct ggml_tensor * src1 = dst->src[1];

GGML_TENSOR_BINARY_OP_LOCALS

const bool src1_cont = ggml_is_contiguous(src1);

ggml_vec_dot_t const vec_dot      = type_traits_cpu[type].vec_dot;
enum ggml_type const vec_dot_type = type_traits_cpu[type].vec_dot_type;

// broadcast factors
const int64_t r2 = ne12 / ne02;
const int64_t r3 = ne13 / ne03;

//printf("ir0_start = %6lld, ir0_end = %6lld, ir1_start = %6lld, ir1_end = %6lld\n", ir0_start, ir0_end, ir1_start, ir1_end);

// threads with no work simply yield (not sure if it helps)
if (ir0_start >= ir0_end || ir1_start >= ir1_end) {
    return;
}

const void * wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata;
const size_t row_size = ggml_row_size(vec_dot_type, ne10);

assert(ne12 % ne02 == 0);
assert(ne13 % ne03 == 0);

// block-tiling attempt
const int64_t blck_0 = 4;
const int64_t blck_1 = 29;

const size_t src1_col_stride = src1_cont || src1->type != vec_dot_type ? row_size : nb11;

// attempt to reduce false-sharing (does not seem to make a difference)
// 16 * 2, accounting for mmla kernels
float tmp[blck_0 * blck_1];

for (int64_t iir1 = ir1_start; iir1 < ir1_end; iir1 += blck_1) {
    for (int64_t iir0 = ir0_start; iir0 < ir0_end; iir0 += blck_0) {
        for (int64_t ir1 = iir1; ir1 < iir1 + blck_1 && ir1 < ir1_end; ir1 += num_rows_per_vec_dot) {
            const int64_t i13 = (ir1 / (ne12 * ne1));
            const int64_t i12 = (ir1 - i13 * ne12 * ne1) / ne1;
            const int64_t i11 = (ir1 - i13 * ne12 * ne1 - i12 * ne1);

            // broadcast src0 into src1
            const int64_t i03 = i13 / r3;
            const int64_t i02 = i12 / r2;

            const int64_t i1 = i11;
            const int64_t i2 = i12;
            const int64_t i3 = i13;

            const char * src0_row = (const char*)src0->data + (0 + i02 * nb02 + i03 * nb03);

            // desc: when src1 is not a contiguous memory block we have to calculate the offset using the strides
            //       if it is, then we have either copied the data to params->wdata and made it contiguous or we are using
            //       the original src1 data pointer, so we should index using the indices directly
            // TODO: this is a bit of a hack, we should probably have a better way to handle this
            const char * src1_col = (const char*)wdata +
                (src1_cont || src1->type != vec_dot_type
                    ? (i11 + i12 * ne11 + i13 * ne12 * ne11) * row_size
                    : (i11 * nb11 + i12 * nb12 + i13 * nb13));
            float * dst_col = (float*)((char*)dst->data + (i1 * nb1 + i2 * nb2 + i3 * nb3));

            //for (int64_t ir0 = iir0; ir0 < iir0 + blck_0 && ir0 < ir0_end; ++ir0) {
            //    vec_dot(ne00, &dst_col[ir0], src0_row + ir0*nb01, src1_col);
            //}

            for (int64_t ir0 = iir0; ir0 < iir0 + blck_0 && ir0 < ir0_end; ir0 += num_rows_per_vec_dot) {
                vec_dot(ne00, &tmp[ir0 - iir0], (num_rows_per_vec_dot > 1 ? blck_0 : 0), src0_row + ir0 * nb01, (num_rows_per_vec_dot > 1 ? nb01 : 0), src1_col, (num_rows_per_vec_dot > 1 ? src1_col_stride : 0), num_rows_per_vec_dot);
            }

            for (int cn = 0; cn < num_rows_per_vec_dot; ++cn) {
                memcpy(&dst_col[iir0 + cn * nb1 / nb0], tmp + (cn * blck_1), (MIN(iir0 + blck_0, ir0_end) - iir0) * sizeof(float));
            }
        }
    }
}

}

試したこと

ここに問題に対して試したことを記載してください。

補足情報（FW/ツールのバージョンなど）

ここにより詳細な情報を記載してください。

行動規範の内容に同意します

回答2件

（以下の回答はあくまで一般論なので実際のところどうなのかについては正確性が保証できません）
まず前提として、コンパイラが色々な最適化を行っているため、Cのソースコードからレジスタが何本使われているかを予測するのは非常に困難です。forループ内なら同じレジスタが使われているという前提も成り立ちません。実際に(gccであれば-Sをつけるなどして)アセンブリのレベルにしてから該当箇所で使われているSIMDレジスタを数えたほうが早いと思います。

つぎに貼っていただいている部分はM*d行列とd*N行列を計算するわけですが、そこに書かれているforループとかtilingブロックとかはMとかNの次元上の話で、dの次元の計算はvec_dotで行われています。つまりvec_dotの内部でSIMDレジスタが大量に使われています。そのためtmp[blck_0 * blck_1]の大きさはSIMDの本数とはあまり関係がありません（あるかもしれません。本当のところはコンパイルしないと分かりません）

基本的にTilingに言及するときは、そのタイル全体がキャッシュ（≠レジスタ）にのるかが重要です。そのタイルサイズがキャッシュサイズに収まっていれば多分うまくいきますが、本当にうまくいっているかどうかについては面倒なことにアセンブリからも分かりません（キャッシュに載っていようが載っていなかろうがメモリアクセスとして表現されるので）。速くなってれば載ってるんだろうなあくらいで考えるか、もしくはIntel VTune Profilerを使うとかですかね...

cache周りの調査を通してプログラムのチューニングしているブログ記事を参考に貼っておきます
https://netflixtechblog.com/seeing-through-hardware-counters-a-journey-to-threefold-performance-increase-2721924a2822

投稿2025/02/20 09:52