Shared embedding kernel

torchao/experimental/kernels/cpu/aarch64/embedding/embedding.h

@@ -10,8 +10,10 @@
 
 #include <arm_neon.h>
 #include <torchao/experimental/kernels/cpu/aarch64/bitpacking/bitpack.h>
+#include <torchao/experimental/kernels/cpu/aarch64/linear/pack_weights.h>
 #include <torchao/experimental/kernels/cpu/aarch64/macro.h>
 #include <cassert>
+#include <vector>
 
 namespace torchao::kernels::cpu::aarch64::embedding {
 
@@ -177,6 +179,7 @@ inline void embedding_(
     if (weight_zeros != nullptr) {
       zero = weight_zeros[group_idx];
     }
+
     internal::vec_dequantize_and_store_16_values(out + i, qvals0, scale, zero);
     internal::vec_dequantize_and_store_16_values(
         out + i + 16, qvals1, scale, zero);
@@ -322,6 +325,61 @@ inline void pack_embedding_weight_qvals(
       qvals + index * embedding_dim);
 }
 
+// Embedding op that shares weights with unembedding linear op
+template <int weight_nbit, int nr, int kr, int sr>
+inline void shared_embedding(
+    // Output
+    float* out,
+    // Inputs
+    const void* packed_weights,
+    int n,
+    int k,
+    int group_size,
+    bool has_weight_zeros,
+    bool has_bias,
+    int index) {
+  assert(k % group_size == 0);
+  assert(group_size % 16 == 0);
+
+  int groups_per_k = k / group_size;
+  std::vector<int8_t> weight_qvals(k * nr);
+  std::vector<float> weight_scales(groups_per_k * nr);
+  std::vector<int8_t> weight_zeros(groups_per_k * nr);
+  std::vector<float> bias(nr);
+
+  // Set n_idx to multiple of nr that is at most index
+  // j is index of "index" in nr group
+  int n_idx = index / nr;
+  n_idx = n_idx * nr;
+  int j = index - n_idx;
+
+  torchao::kernels::cpu::aarch64::linear::packing::
+      unpack_weights_at_n_idx<weight_nbit, nr, kr, sr>(
+          weight_qvals.data(),
+          weight_scales.data(),
+          has_weight_zeros ? weight_zeros.data() : nullptr,
+          has_bias ? bias.data() : nullptr,
+          n_idx,
+          n,
+          k,
+          group_size,
+          has_weight_zeros,
+          has_bias,
+          packed_weights);
+
+  // Dequantize and store to output (size k)
+  int8x16_t qvals;
+  for (int i = 0; i < k; i += 16) {
+    qvals = vld1q_s8(weight_qvals.data() + j * k + i);
+    float scale = weight_scales[j * groups_per_k + i / group_size];
+    float zero = 0.0;
+    if (has_weight_zeros) {
+      zero = weight_zeros[j * groups_per_k + i / group_size];
+    }
+    internal::vec_dequantize_and_store_16_values(out + i, qvals, scale, zero);
+  }
+}
+
 } // namespace torchao::kernels::cpu::aarch64::embedding
 
 #endif // defined(__aarch64__) || defined(__ARM_NEON)

torchao/experimental/kernels/cpu/aarch64/linear/pack_weights.h

@@ -339,7 +339,7 @@ void unpack_weights_at_n_idx(
     int group_size,
     bool has_weight_zeros,
     bool has_bias,
-    void* packed_weights) {
+    const void* packed_weights) {
   assert(k % group_size == 0);
   assert(group_size % kr == 0);
   assert(n_idx % nr == 0);
@@ -441,7 +441,7 @@ void unpack_weights(
     int group_size,
     bool has_weight_zeros,
     bool has_bias,
-    void* packed_weights) {
+    const void* packed_weights) {
   assert(k % group_size == 0);
   assert(group_size % kr == 0);
   int groups_per_k = k / group_size;

torchao/experimental/kernels/cpu/aarch64/tests/test_embedding.cpp

@@ -8,6 +8,7 @@
 
 #include <gtest/gtest.h>
 #include <torchao/experimental/kernels/cpu/aarch64/embedding/embedding.h>
+#include <torchao/experimental/kernels/cpu/aarch64/linear/pack_weights.h>
 #include <torchao/experimental/kernels/cpu/aarch64/tests/test_utils.h>
 #include <vector>
 
@@ -53,6 +54,54 @@ void test_embedding(
   }
 }
 
+template <int weight_nbit, int nr, int kr, int sr>
+void test_shared_embedding(
+    int num_embeddings,
+    int embedding_dim,
+    int group_size,
+    bool has_weight_zeros) {
+  auto test_case = torchao::lowbit_embedding_test_case<weight_nbit>::generate(
+      num_embeddings, embedding_dim, group_size, has_weight_zeros);
+
+  // Pack weights for linear op
+  int n = num_embeddings;
+  int k = embedding_dim;
+  bool has_bias = false;
+  float* bias = nullptr;
+  std::vector<char> packed_weights(
+      torchao::kernels::cpu::aarch64::linear::packing::packed_weights_size(
+          n, k, group_size, weight_nbit, has_weight_zeros, has_bias, nr));
+  torchao::kernels::cpu::aarch64::linear::packing::
+      pack_weights<weight_nbit, nr, kr, sr>(
+          packed_weights.data(),
+          n,
+          k,
+          group_size,
+          test_case.weight_qvals.data(),
+          test_case.weight_scales.data(),
+          has_weight_zeros ? test_case.weight_zeros.data() : nullptr,
+          bias);
+
+  // Call shared_embedding
+  auto output = std::vector<float>(num_embeddings * embedding_dim, 0.0);
+  for (int i = 0; i < num_embeddings; i++) {
+    torchao::kernels::cpu::aarch64::embedding::
+        shared_embedding<weight_nbit, nr, kr, sr>(
+            output.data() + i * embedding_dim,
+            packed_weights.data(),
+            n,
+            k,
+            group_size,
+            has_weight_zeros,
+            has_bias,
+            i);
+  }
+
+  for (int i = 0; i < num_embeddings * embedding_dim; i++) {
+    EXPECT_NEAR(output[i], test_case.expected_outputs[i], kTol);
+  }
+}
+
 TEST(test_embedding, NBit1) {
   constexpr int num_embeddings = 5;
   constexpr int group_size = 128 * 3 + 64 + 32;
@@ -97,7 +146,7 @@ TEST(test_embedding, NBit4) {
       num_embeddings, embedding_dim, group_size, /*has_weight_zeros=*/false);
 
   // More detailed testing for 4-bit case
-  
+
   test_embedding<4>(
       num_embeddings,
       /*embedding_dim=*/256,
@@ -152,4 +201,14 @@ TEST(test_embedding, NBit6) {
       num_embeddings, embedding_dim, group_size, /*has_weight_zeros=*/false);
 }
 
+TEST(test_embedding, SharedEmbeddingTest) {
+  constexpr int weight_nbit = 3;
+  constexpr int num_embeddings = 17;
+  constexpr int group_size = 64;
+  constexpr int embedding_dim = group_size * 7;
+
+  test_shared_embedding<weight_nbit, /*nr*/ 8, /*kr*/ 16, /*sr*/ 2>(
+      num_embeddings, embedding_dim, group_size, /*has_weight_zeros*/true);
+}
+
 #endif // defined(__aarch64__) || defined(__ARM_NEON)