Skip to content

Commit 8f6596d

Browse files
authored
Port signal custom ops to TensorFlow and TFLite C APIs (#3880)
* Port signal custom ops to TensorFlow C API and restore BUILD files Vendor the 16 pure-C TensorFlow headers in third_party/tf_c_headers and port all 9 signal custom ops and kernels from tensorflow/core/framework C++ headers to tensorflow/c/{ops,kernels}.h. Restore the Bazel BUILD files and tflm_signal.bzl so the signal custom-op shared library plugins build without downloading TensorFlow at build time. BUG=n/a * Add TFLite C API kernels for signal custom ops Vendor the 12 pure-C TFLite Opaque API headers in third_party/tflite_c_headers and implement all 14 signal custom ops in signal/lite/kernels using TfLiteOperator. Each op exposes Create*Operator() for LiteRT/C API callers and Register_*() for TFLite MutableOpResolver callers, without requiring a TFLite repo dependency at build time. BUG=n/a * Fix clang-tidy warnings in signal C API ops and kernels - Move GetFlatSize and MakeRegistration to kernel_util.cc so compile_commands.json includes -Ithird_party/tflite_c_headers for kernel_util.h. - Replace static lambda initializers in signal/tensorflow_core with __attribute__((constructor)) functions. - Cast integer operands before multiplication and replace long long / signed bitwise shift. BUG=n/a * Simplify signal C API ops and kernels with C++ RAII wrappers Wrap TensorFlow and TFLite C API handles, status propagation, and registration callbacks in small C++ RAII helpers so ops and kernels read like standard C++ classes without manual memory management. BUG=n/a * Add sync scripts and update READMEs for C API headers Add sync.sh to third_party/tf_c_headers and third_party/tflite_c_headers to fetch the pure-C headers verbatim from upstream TensorFlow (pinned to v2.20.0 by default, or from a local checkout). Document how to run the sync scripts in both README.md files. BUG=n/a * Add TFLite and LiteRT registration and tests for signal custom ops * Consolidate signal build rules in py_rules.bzl and add tflm_oss_manual_tags * Rename OpDefBuilder to SignalOpDefBuilder and document signal ops in Python - Rename tflite::tflm_signal::OpDefBuilder to SignalOpDefBuilder in shape_inference_utils.h and clarify that it wraps TF_OpDefinitionBuilder* from tensorflow/c/ops.h (which has no .Doc() API). - Replace detached REGISTER_OP doc blocks in signal/tensorflow_core/ops/*.cc with uniform registration comments pointing to the Python op modules. - Replace internal _*_wrapper closures in python/tflite_micro/signal/ops/*.py with direct top-level functions and complete Google-style docstrings on all 14 public signal ops. BUG=n/a
1 parent f7ed40d commit 8f6596d

100 files changed

Lines changed: 15523 additions & 1485 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎python/py_rules.bzl‎

Lines changed: 125 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,17 +17,39 @@ target under the tflite_micro namespace.
1717
"""
1818

1919
load("@pybind11_bazel//:build_defs.bzl", _pybind_extension = "pybind_extension", _pybind_library = "pybind_library")
20+
load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
21+
load("@rules_cc//cc:cc_library.bzl", "cc_library")
2022
load("@rules_python//python:defs.bzl", "py_binary", "py_library", "py_test")
2123
load("@rules_python//python:packaging.bzl", _py_package = "py_package", _py_wheel = "py_wheel")
2224
load("@rules_python//python:pip.bzl", _compile_pip_requirements = "compile_pip_requirements")
2325
load("@rules_shell//shell:sh_test.bzl", "sh_test")
26+
load("@tflm_pip_deps//:requirements.bzl", "requirement")
27+
load("//tensorflow/lite/micro:build_def.bzl", "INCOMPATIBLE_WITH_WINDOWS")
2428

2529
compile_pip_requirements = _compile_pip_requirements
2630
py_package = _py_package
2731
py_wheel = _py_wheel
28-
pybind_extension = _pybind_extension
2932
pybind_library = _pybind_library
3033

34+
def tflm_oss_manual_tags():
35+
"""Returns ["manual"] in OSS builds for targets requiring external TF/LiteRT deps."""
36+
return ["manual"]
37+
38+
def pybind_extension(name, **kwargs):
39+
if "manual" in kwargs.get("tags", []):
40+
_pybind_library(
41+
name = name + "_pybind_lib",
42+
**kwargs
43+
)
44+
cc_binary(
45+
name = name + ".so",
46+
linkshared = True,
47+
tags = kwargs.get("tags", []),
48+
deps = [":" + name + "_pybind_lib"],
49+
)
50+
else:
51+
_pybind_extension(name = name, **kwargs)
52+
3153
def tflm_py_library(deps = [], **kwargs):
3254
py_library(
3355
deps = deps + [
@@ -75,3 +97,105 @@ def tflm_whl_test(name, srcs, **kwargs):
7597
],
7698
**kwargs
7799
)
100+
101+
def py_tflm_signal_library(
102+
name,
103+
srcs = [],
104+
deps = [],
105+
visibility = None,
106+
cc_op_defs = [],
107+
cc_op_kernels = []):
108+
"""Creates build rules for signal custom ops as shared libraries."""
109+
binary_path = "ops"
110+
if srcs:
111+
binary_path_end_pos = srcs[0].rfind("/")
112+
binary_path = srcs[0][0:binary_path_end_pos]
113+
114+
binary_name = binary_path + "/_" + (cc_op_kernels[0][1:] if not cc_op_defs else name) + ".so"
115+
116+
if cc_op_defs:
117+
library_name = name + "_cc"
118+
cc_library(
119+
name = library_name,
120+
alwayslink = 1,
121+
copts = ["-pthread"],
122+
target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
123+
deps = cc_op_defs + cc_op_kernels,
124+
)
125+
cc_binary(
126+
name = binary_name,
127+
copts = ["-pthread"],
128+
linkopts = select({
129+
"@platforms//os:macos": ["-Wl,-undefined,dynamic_lookup"],
130+
"//conditions:default": ["-Wl,--unresolved-symbols=ignore-in-object-files"],
131+
}),
132+
linkshared = 1,
133+
target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
134+
deps = [
135+
":" + library_name,
136+
],
137+
)
138+
139+
tflm_py_library(
140+
name = name,
141+
srcs = srcs,
142+
data = [":" + binary_name] if cc_op_defs else [],
143+
target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
144+
visibility = visibility,
145+
deps = deps,
146+
)
147+
148+
def tflm_signal_op_test(
149+
name,
150+
srcs,
151+
data = [],
152+
deps = [],
153+
tags = [],
154+
target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
155+
**kwargs):
156+
"""Creates TFLM and TFLite test targets for a signal custom op."""
157+
common_tags = tags + tflm_oss_manual_tags() + [
158+
"noasan",
159+
"nomsan",
160+
"noubsan",
161+
]
162+
tflm_py_test(
163+
name = name,
164+
srcs = srcs,
165+
data = data,
166+
tags = common_tags,
167+
target_compatible_with = target_compatible_with,
168+
deps = deps,
169+
**kwargs
170+
)
171+
tflm_py_test(
172+
name = name + "_tflite",
173+
srcs = srcs,
174+
main = srcs[0],
175+
data = data + [
176+
"//signal/lite/register:pywrap_signal_ops.so",
177+
],
178+
env = {"TFLM_SIGNAL_USE_TFLITE": "1"},
179+
tags = common_tags,
180+
target_compatible_with = target_compatible_with,
181+
deps = deps + [
182+
requirement("ai-edge-litert"),
183+
],
184+
**kwargs
185+
)
186+
187+
def tflm_signal_kernel_library(
188+
name,
189+
srcs = [],
190+
hdrs = [],
191+
deps = [],
192+
copts = [],
193+
alwayslink = 1):
194+
cc_library(
195+
name = name,
196+
srcs = srcs,
197+
hdrs = hdrs,
198+
alwayslink = alwayslink,
199+
copts = copts,
200+
deps = deps,
201+
)

0 commit comments

Comments
 (0)