executorch/backends/vulkan/test/custom_ops/targets.bzl at 885d27fc97d707f6d16138fe01e36175170d7edc · 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
load("@fbsource//tools/build_defs:platform_defs.bzl", "ANDROID")
load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
load(
    "@fbsource//xplat/executorch/backends/vulkan:targets.bzl",
    "get_platforms",
    "vulkan_spv_shader_lib",
)

def define_custom_op_test_binary(custom_op_name, extra_deps = [], src_file = None):
    deps_list = [
        ":prototyping_utils",
        ":operator_implementations",
        ":custom_ops_shaderlib",
        "//executorch/backends/vulkan:vulkan_graph_runtime",
        runtime.external_dep_location("libtorch"),
    ] + extra_deps

    src_file_str = src_file if src_file else "{}.cpp".format(custom_op_name)

    runtime.cxx_binary(
        name = custom_op_name,
        srcs = [
            src_file_str,
        ],
        platforms = get_platforms(),
        define_static_target = False,
        deps = deps_list,
    )

def define_common_targets(is_fbcode = False):
    if is_fbcode:
        return

    # Shader library from GLSL files
    runtime.filegroup(
        name = "custom_ops_shaders",
        srcs = native.glob([
            "glsl/*.glsl",
            "glsl/*.yaml",
        ]),
        visibility = ["PUBLIC"],
    )

    vulkan_spv_shader_lib(
        name = "custom_ops_shaderlib",
        spv_filegroups = {
            ":custom_ops_shaders": "glsl",
        },
        is_fbcode = is_fbcode,
    )

    # Prototyping utilities library
    runtime.cxx_library(
        name = "prototyping_utils",
        srcs = [
            "utils.cpp",
            "conv2d_utils.cpp",
        ],
        headers = [
            "utils.h",
            "conv2d_utils.h",
        ],
        exported_headers = [
            "utils.h",
            "conv2d_utils.h",
        ],
        platforms = get_platforms(),
        deps = [
            "//executorch/backends/vulkan:vulkan_graph_runtime",
        ],
        visibility = ["PUBLIC"],
    )

    # Operator implementations library
    runtime.cxx_library(
        name = "operator_implementations",
        srcs = native.glob([
            "impl/*.cpp",
        ]),
        platforms = get_platforms(),
        deps = [
            "//executorch/backends/vulkan:vulkan_graph_runtime",
            ":custom_ops_shaderlib",
        ],
        visibility = ["PUBLIC"],
        link_whole = True,
    )

    define_custom_op_test_binary("add")
    define_custom_op_test_binary("q8csw_linear")
    define_custom_op_test_binary("q8csw_conv2d")
    define_custom_op_test_binary("choose_qparams_per_row")
    define_custom_op_test_binary("q4gsw_linear")
    define_custom_op_test_binary("qdq8ta_conv2d_activations")
    define_custom_op_test_binary("q8ta_q8csw_q8to_conv2d")
    define_custom_op_test_binary("q8ta_q8csw_q8to_conv2d_dw")
    define_custom_op_test_binary("q8ta_q8ta_q8to_add")