[kernel][mobile] add mobile-optimized CPU backend

[AlpinDale]

WIP, only targeting Android for now.

[gemini-code-assist]

Code Review

This pull request introduces a mobile-optimized CPU backend targeting Android, leveraging ARM NEON for performance. The changes are extensive, adding numerous C++ files with optimized kernels for various operations like attention, matrix multiplication, and activations, along with the necessary CMake build system modifications. While the effort to optimize for mobile is commendable, the review has identified several critical issues. These include correctness bugs such as potential race conditions due to thread-unsafe static variables, incorrect arithmetic leading to wrong results, and linker errors from function name mismatches. There are also significant performance issues stemming from inefficient memory access patterns and repeated quantization/dequantization operations within tight loops. Addressing these critical and high-severity issues is essential for the stability and performance of the new mobile backend.

[gemini-code-assist]

The accumulation of attention-weighted values into the output vector o_vec is inefficient and can lead to precision loss. Inside the kv_pos loop, the o_vec (which is int8_t) is repeatedly dequantized to float, updated, and then requantized. This process is computationally expensive and introduces quantization errors at each step.

A more efficient and accurate approach is to use a temporary floating-point accumulator for the output vector. This accumulator can sum up all the weighted values in full precision. The final quantization to int8_t should only happen once, after the loop over kv_pos is complete.

              std::vector<float> o_vec_fp32(head_dim, 0.0f);

              for (size_t kv_pos = 0; kv_pos < kv_seq_len; ++kv_pos) {
                const float attn_weight = scores_row[kv_pos];
                if (attn_weight == 0.0f) continue;

                const int8_t* v_vec =
                    V_base + kv_pos * kv_seq_stride + kv_head_idx * head_dim;

                for (size_t dim = 0; dim < head_dim; ++dim) {
                  o_vec_fp32[dim] +=
                      attn_weight * static_cast<float>(v_vec[dim]) * v_scale;
                }
              }

              int8_t* o_vec =
                  O_base + q_pos * o_seq_stride + q_head_idx * head_dim;
              for (size_t dim = 0; dim < head_dim; ++dim) {
                float val = o_vec_fp32[dim];
                int32_t quantized_result =
                    static_cast<int32_t>(val / output_scale +
                                         (val >= 0 ? 0.5f : -0.5f));
                quantized_result =
                    std::max(-128, std::min(127, quantized_result));
                o_vec[dim] = static_cast<int8_t>(quantized_result);
              }

[gemini-code-assist]

The vectorization logic in rms_norm_i8_f32 is inefficient and likely incorrect due to overlapping memory loads. The loop iterates with a step of TILE_SIZE (16), and the inner loop with UNROLL_FACTOR (4) and SIMD_WIDTH (4) results in addresses i, i+4, i+8, i+12. At each of these addresses, vld1_s8 loads 8 bytes. This means, for example, the second load at i+4 re-loads 4 bytes that were already loaded at i. Furthermore, from each 8-byte load, only the first 4 bytes are processed (vget_low_s16(vmovl_s8(input_i8))), discarding half of the loaded data.

This should be corrected to avoid overlapping loads and to process all loaded data, which will significantly improve performance and ensure correctness. A similar issue exists in the second part of the function where the normalization is applied.

[gemini-code-assist]

There is an integer division bug in the calculation of freq. The variables i and head_dim are of type size_t, so (2.0f * i) / head_dim will perform integer division before the multiplication, leading to incorrect frequency values. You should cast these variables to float to ensure floating-point division.

Suggested change

	const float freq = 1.0f / powf(theta, (2.0f * i) / head_dim);
	const float freq = 1.0f / powf(theta, (2.0f * static_cast<float>(i)) / static_cast<float>(head_dim));

[gemini-code-assist]

The logic for handling the remainder of the K dimension is incorrect. The scalar dot_product is broadcast to a vector and added to the vector accumulator accumulators[m][n]. When vaddvq_s32 is called later to get the final sum, the dot_product is effectively added four times, leading to incorrect matrix multiplication results.

The correct approach is to perform the horizontal sum of the vector accumulator first to get a scalar sum, and then add the scalar-computed remainder to it.

      for (int m = 0; m < TILE_M; ++m) {
        size_t row = row_block + m;
        if (row >= M) continue;
        for (int n = 0; n < TILE_N; ++n) {
          size_t col = col_block + n;
          if (col >= N) continue;
          int32_t sum = vaddvq_s32(accumulators[m][n]);
          for (size_t k = K_aligned; k < K; ++k) {
            sum += static_cast<int32_t>(a[row * K + k]) *
                   static_cast<int32_t>(b_transposed[col * K + k]);
          }
          c[row * N + col] = sum;
        }
      }

[gemini-code-assist]

The use of a static variable token_history makes the sample_f16 function not thread-safe. If this function is called from multiple threads concurrently (e.g., for different sequences in a batch), they will all share and modify the same token_history vector, leading to race conditions and incorrect behavior.

To fix this, token_history should be managed externally and passed into the function as an argument. The related constants MAX_HISTORY and REPETITION_PENALTY should also be passed as arguments if they need to be configurable per call.

[gemini-code-assist]

The code for storing the transposed columns (col0, col1, col2, col3) is highly repetitive. This duplication makes the code harder to read and maintain. You can refactor this section by using a helper lambda or a small loop to handle the storage of each column, which would significantly reduce code duplication and improve clarity.