executorch/kernels/quantized/cpu/targets.bzl at 444ee9db98268ea9fbdd14866c9edcb51e2777d7 · pytorch/executorch · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
load("@fbsource//xplat/executorch/kernels/portable:op_registration_util.bzl", "define_op_target", "op_target")

_QUANT_OPS = (
    op_target(
        name = "op_add",
    ),
    op_target(
        name = "op_choose_qparams",
        deps = [
            "//executorch/kernels/portable/cpu:vec_ops",
            "//executorch/extension/threadpool:threadpool",
        ],
    ),
    op_target(
        name = "op_dequantize",
        deps = [
            "//executorch/kernels/portable/cpu/util:reduce_util",
        ],
        _aten_mode_deps = [
            "//executorch/kernels/portable/cpu/util:reduce_util_aten",
        ],
    ),
    op_target(
        name = "op_embedding",
    ),
    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",
        deps = [
            "//executorch/kernels/portable/cpu:vec_ops",
        ],
    ),
    op_target(
        name = "op_mixed_linear",
        deps = [
            "//executorch/kernels/portable/cpu:vec_ops",
        ],
    ),
    op_target(
        name = "op_quantize",
    ),
)

def define_common_targets():
    for op in _QUANT_OPS:
        define_op_target(is_aten_op = False, **op)

    quant_op_targets = [":{}".format(op["name"]) for op in _QUANT_OPS]

    runtime.cxx_library(
        name = "quantized_cpu",
        srcs = [],
        visibility = [
            "//executorch/kernels/quantized/...",
            "//executorch/extension/pybindings/test/...",
        ],
        exported_deps = quant_op_targets,
    )

    runtime.cxx_library(
        name = "embeddingxb",
        srcs = ["embeddingxb.cpp"],
        exported_headers = ["embeddingxb.h"],
        visibility = [
            "//executorch/kernels/quantized/...",
        ],
        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 = [],
        visibility = ["//executorch/kernels/quantized/..."],
        exported_deps = [t + "_aten" for t in quant_op_targets],
    )