• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1load(
2    "//tensorflow:tensorflow.bzl",
3    "tf_binary_additional_srcs",
4)
5
6# Generate Java wrapper classes for all registered core operations and package
7# them into a single source archive (.srcjar).
8#
9# For example:
10#  tf_java_op_gen_srcjar("gen_sources", ":gen_tool", "my.package")
11#
12# will create a genrule named "gen_sources" that generates source files under
13#     ops/src/main/java/my/package/**/*.java
14#
15# and then archive those source files into
16#     ops/gen_sources.srcjar
17#
18def tf_java_op_gen_srcjar(
19        name,
20        gen_tool,
21        base_package,
22        api_def_srcs = [],
23        out_dir = "ops/",
24        out_src_dir = "src/main/java/",
25        visibility = ["//tensorflow/java:__pkg__"]):
26    gen_cmds = ["rm -rf $(@D)"]  # Always start from fresh when generating source files
27    srcs = api_def_srcs[:]
28
29    if not api_def_srcs:
30        api_def_args_str = ","
31    else:
32        api_def_args = []
33        for api_def_src in api_def_srcs:
34            # Add directory of the first ApiDef source to args.
35            # We are assuming all ApiDefs in a single api_def_src are in the
36            # same directory.
37            api_def_args.append(
38                "$$(dirname $$(echo $(locations " + api_def_src +
39                ") | cut -d\" \" -f1))",
40            )
41        api_def_args_str = ",".join(api_def_args)
42
43    gen_cmds += ["$(location " + gen_tool + ")" +
44                 " --output_dir=$(@D)/" + out_src_dir +
45                 " --base_package=" + base_package +
46                 " --api_dirs=" + api_def_args_str]
47
48    # Generate a source archive containing generated code for these ops.
49    gen_srcjar = out_dir + name + ".srcjar"
50    gen_cmds += ["$(JAVABASE)/bin/jar cMf $(location :" + gen_srcjar + ") -C $(@D) src"]
51
52    native.genrule(
53        name = name,
54        srcs = srcs,
55        outs = [gen_srcjar],
56        tools = [
57            # copybara:uncomment_begin(using system-provided in OSS build)
58            # "//third_party/java/jdk/jar:jar",
59            # "//third_party/java/jdk:jdk",
60            # copybara:uncomment_end
61            gen_tool,
62        ] + tf_binary_additional_srcs(),
63        toolchains = ["@bazel_tools//tools/jdk:current_host_java_runtime"],
64        cmd = " && ".join(gen_cmds),
65    )
66