xref: /aosp_15_r20/external/executorch/kernels/quantized/cpu/targets.bzl (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
1load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
2load("@fbsource//xplat/executorch/kernels/portable:op_registration_util.bzl", "define_op_target", "op_target")
3
4_QUANT_OPS = (
5    op_target(
6        name = "op_add",
7    ),
8    op_target(
9        name = "op_choose_qparams",
10        deps = [
11            "//executorch/kernels/portable/cpu:vec_ops",
12        ],
13    ),
14    op_target(
15        name = "op_dequantize",
16        deps = [
17            "//executorch/kernels/portable/cpu/util:reduce_util",
18        ],
19        _aten_mode_deps = [
20            "//executorch/kernels/portable/cpu/util:reduce_util_aten",
21        ],
22    ),
23    op_target(
24        name = "op_embedding",
25    ),
26    op_target(
27        name = "op_embedding2b",
28        deps = ["//executorch/kernels/quantized/cpu:embeddingxb"],
29        _aten_mode_deps = [
30            "//executorch/kernels/quantized/cpu:embeddingxb_aten",
31        ],
32    ),
33    op_target(
34        name = "op_embedding4b",
35        deps = ["//executorch/kernels/quantized/cpu:embeddingxb"],
36         _aten_mode_deps = [
37            "//executorch/kernels/quantized/cpu:embeddingxb_aten",
38        ],
39    ),
40    op_target(
41        name = "op_mixed_mm",
42        deps = [
43            "//executorch/kernels/portable/cpu:vec_ops",
44        ],
45    ),
46    op_target(
47        name = "op_mixed_linear",
48        deps = [
49            "//executorch/kernels/portable/cpu:vec_ops",
50        ],
51    ),
52    op_target(
53        name = "op_quantize",
54        deps = [
55            "//executorch/kernels/portable/cpu/util:reduce_util",
56        ],
57        _aten_mode_deps = [
58            "//executorch/kernels/portable/cpu/util:reduce_util_aten",
59        ],
60    ),
61)
62
63def define_common_targets():
64    for op in _QUANT_OPS:
65        define_op_target(is_aten_op = False, **op)
66
67    quant_op_targets = [":{}".format(op["name"]) for op in _QUANT_OPS]
68
69    runtime.cxx_library(
70        name = "quantized_cpu",
71        srcs = [],
72        visibility = [
73            "//executorch/kernels/quantized/...",
74            "//executorch/extension/pybindings/test/...",
75        ],
76        exported_deps = quant_op_targets,
77    )
78
79    runtime.cxx_library(
80        name = "embeddingxb",
81        srcs = ["embeddingxb.cpp"],
82        exported_headers = ["embeddingxb.h"],
83        visibility = [
84            "//executorch/kernels/quantized/...",
85        ],
86        deps = ["//executorch/runtime/kernel:kernel_includes"],
87    )
88
89    runtime.cxx_library(
90        name = "embeddingxb_aten",
91        srcs = ["embeddingxb.cpp"],
92        exported_headers = ["embeddingxb.h"],
93        visibility = [
94            "//executorch/kernels/quantized/...",
95        ],
96        deps = ["//executorch/runtime/kernel:kernel_includes_aten"],
97    )
98
99    runtime.cxx_library(
100        name = "quantized_cpu_aten",
101        srcs = [],
102        visibility = ["//executorch/kernels/quantized/..."],
103        exported_deps = [t + "_aten" for t in quant_op_targets],
104    )
105