add aten target to bazel

Author: metascroy

kernels/quantized/cpu/op_embedding2b.cpp

@@ -74,6 +74,27 @@ Tensor& quantized_embedding_2bit_out(
       2);
 }
 
+Tensor& quantized_embedding_2bit_dtype_out(
+    const Tensor& weight,
+    const Tensor& weight_scales,
+    const optional<Tensor>& opt_weight_zero_points,
+    int64_t weight_quant_min,
+    int64_t weight_quant_max,
+    const Tensor& indices,
+    exec_aten::optional<ScalarType> out_dtype,
+    Tensor& out) {
+  return quantized_embedding_xbit_dtype_out(
+      weight,
+      weight_scales,
+      opt_weight_zero_points,
+      weight_quant_min,
+      weight_quant_max,
+      indices,
+      out_dtype,
+      out,
+      2);
+}
+
 Tensor& quantized_embedding_2bit_dtype_out(
     KernelRuntimeContext& context,
     const Tensor& weight,

kernels/quantized/cpu/targets.bzl

@@ -26,10 +26,16 @@ _QUANT_OPS = (
     op_target(
         name = "op_embedding2b",
         deps = ["//executorch/kernels/quantized/cpu:embeddingxb"],
+        _aten_mode_deps = [
+            "//executorch/kernels/quantized/cpu:embeddingxb_aten",
+        ],
     ),
     op_target(
         name = "op_embedding4b",
         deps = ["//executorch/kernels/quantized/cpu:embeddingxb"],
+         _aten_mode_deps = [
+            "//executorch/kernels/quantized/cpu:embeddingxb_aten",
+        ],
     ),
     op_target(
         name = "op_mixed_mm",
@@ -80,6 +86,16 @@ def define_common_targets():
         deps = ["//executorch/runtime/kernel:kernel_includes"],
     )
 
+    runtime.cxx_library(
+        name = "embeddingxb_aten",
+        srcs = ["embeddingxb.cpp"],
+        exported_headers = ["embeddingxb.h"],
+        visibility = [
+            "//executorch/kernels/quantized/...",
+        ],
+        deps = ["//executorch/runtime/kernel:kernel_includes_aten"],
+    )
+
     runtime.cxx_library(
         name = "quantized_cpu_aten",
         srcs = [],

kernels/quantized/test/op_embedding2b_test.cpp

@@ -159,3 +159,33 @@ TEST(OpQuantizedEmbedding2bTest, TestGroupWiseQuantizedEmbeddingDeath2) {
           out),
       "");
 }
+
+TEST(OpQuantizedEmbedding2bTest, TestGroupWiseQuantizedEmbeddingDeath3) {
+  et_pal_init();
+  TensorFactory<ScalarType::Byte> tfb;
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> tfl;
+
+  int64_t quant_min = -2;
+  int64_t quant_max = 1;
+
+  Tensor weight_scales =
+      tf.make({2, 3}, {1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0});
+  Tensor weight_zero_points = tf.make({2, 3}, {0, 0, 0, 0, 0, 0, 0, 0, 0});
+  Tensor qweight = tfb.make({2, 1}, {236, 134, 228});
+  Tensor indices = tfl.make({2}, {0, 2});
+  Tensor out = tf.zeros({2, 8});
+
+  // scales/zeros imply groupsize 3, which does not divide embed dimension from
+  // qvals (8)
+  ET_EXPECT_DEATH(
+      quantized_embedding_2bit_out(
+          qweight,
+          weight_scales,
+          weight_zero_points,
+          quant_min,
+          quant_max,
+          indices,
+          out),
+      "");
+}