xref: /aosp_15_r20/external/pytorch/third_party/onnx.BUILD (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1load("@rules_proto//proto:defs.bzl", "proto_library")
2load("@rules_cc//cc:defs.bzl", "cc_library", "cc_proto_library")
3load("@rules_python//python:defs.bzl", "py_binary")
4
5py_binary(
6    name = "gen_proto",
7    srcs = ["onnx/gen_proto.py"],
8    data = [
9        "onnx/onnx.in.proto",
10        "onnx/onnx-operators.in.proto",
11        "onnx/onnx-data.in.proto",
12    ],
13)
14
15genrule(
16    name = "generate_onnx_proto",
17    outs = [
18        "onnx/onnx_onnx_torch-ml.proto",
19        "onnx/onnx-ml.pb.h",
20    ],
21    cmd = "$(location :gen_proto) -p onnx_torch -o $(@D)/onnx onnx -m >/dev/null && sed -i 's/onnx_onnx_torch-ml.pb.h/onnx\\/onnx_onnx_torch-ml.pb.h/g' $(@D)/onnx/onnx-ml.pb.h",
22    tools = [":gen_proto"],
23)
24
25genrule(
26    name = "generate_onnx_operators_proto",
27    outs = [
28        "onnx/onnx-operators_onnx_torch-ml.proto",
29        "onnx/onnx-operators-ml.pb.h",
30    ],
31    cmd = "$(location :gen_proto) -p onnx_torch -o $(@D)/onnx onnx-operators -m >/dev/null && sed -i 's/onnx-operators_onnx_torch-ml.pb.h/onnx\\/onnx-operators_onnx_torch-ml.pb.h/g' $(@D)/onnx/onnx-operators-ml.pb.h",
32    tools = [":gen_proto"],
33)
34
35genrule(
36    name = "generate_onnx_data_proto",
37    outs = [
38        "onnx/onnx-data_onnx_torch.proto",
39        "onnx/onnx-data.pb.h",
40    ],
41    cmd = "$(location :gen_proto) -p onnx_torch -o $(@D)/onnx onnx-data -m >/dev/null && sed -i 's/onnx-data_onnx_torch.pb.h/onnx\\/onnx-data_onnx_torch.pb.h/g' $(@D)/onnx/onnx-data.pb.h",
42    tools = [":gen_proto"],
43)
44
45cc_library(
46    name = "onnx",
47    srcs = glob(
48        [
49            "onnx/*.cc",
50            "onnx/common/*.cc",
51            "onnx/defs/**/*.cc",
52            "onnx/shape_inference/*.cc",
53            "onnx/version_converter/*.cc",
54        ],
55        exclude = [
56            "onnx/cpp2py_export.cc",
57        ],
58    ),
59    hdrs = glob([
60        "onnx/*.h",
61        "onnx/version_converter/*.h",
62        "onnx/common/*.h",
63        "onnx/defs/**/*.h",
64        "onnx/shape_inference/*.h",
65        "onnx/version_converter/adapters/*.h",
66    ]) + [
67        "onnx/onnx-ml.pb.h",
68        "onnx/onnx-operators-ml.pb.h",
69        "onnx/onnx-data.pb.h",
70    ],
71    defines = [
72        "ONNX_ML=1",
73        "ONNX_NAMESPACE=onnx_torch",
74    ],
75    includes = [
76        ".",
77        "onnx/",
78    ],
79    visibility = ["//visibility:public"],
80    deps = [
81        ":onnx_proto_lib",
82    ],
83)
84
85cc_library(
86    name = "onnx_proto_headers",
87    hdrs = glob([
88        "onnx/*_pb.h",
89    ]),
90    visibility = ["//visibility:public"],
91    deps = [
92        ":onnx_proto_lib",
93    ],
94)
95
96proto_library(
97    name = "onnx_proto",
98    srcs = [
99        "onnx/onnx-operators_onnx_torch-ml.proto",
100        "onnx/onnx_onnx_torch-ml.proto",
101        "onnx/onnx-data_onnx_torch.proto",
102    ],
103)
104
105cc_proto_library(
106    name = "onnx_proto_lib",
107    deps = [":onnx_proto"],
108)
109