pytorch

1
load("@rules_proto//proto:defs.bzl", "proto_library")
2
load("@rules_cc//cc:defs.bzl", "cc_library", "cc_proto_library")
3
load("@rules_python//python:defs.bzl", "py_binary")
4

5
py_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

15
genrule(
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

25
genrule(
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

35
genrule(
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

45
cc_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

85
cc_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

96
proto_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

105
cc_proto_library(
106
    name = "onnx_proto_lib",
107
    deps = [":onnx_proto"],
108
)
109