feat: add model compression draft

Add a feature-complete draft of the model compression feature in a single commit. This commit isn't intended for merging into mainline; it exists solely for reviewers and testers to take an early look.
ddavis-2015 · Nov 14, 2024 · 6eac102 · 6eac102
1 parent f91dd91
commit 6eac102
Show file tree

Hide file tree

Showing 138 changed files with 9,571 additions and 1,158 deletions.
diff --git a/BUILD b/BUILD
@@ -7,3 +7,17 @@ refresh_compile_commands(
     name = "refresh_compile_commands",
     targets = ["//..."],
 )
+
+load("@bazel_skylib//rules:common_settings.bzl", "bool_flag")
+
+bool_flag(
+    name = "with_compression",
+    build_setting_default = False,
+)
+
+config_setting(
+    name = "with_compression_enabled",
+    flag_values = {
+        ":with_compression": "True",
+    },
+)
diff --git a/codegen/build_def.bzl b/codegen/build_def.bzl
@@ -1,6 +1,6 @@
 """ Build rule for generating ML inference code from TFLite model. """
 
-load("//tensorflow/lite/micro:build_def.bzl", "micro_copts")
+load("//tensorflow/lite/micro:build_def.bzl", "tflm_cc_library")
 
 def tflm_inference_library(
         name,
@@ -39,7 +39,7 @@ def tflm_inference_library(
         visibility = ["//visibility:private"],
     )
 
-    native.cc_library(
+    tflm_cc_library(
         name = name,
         hdrs = [name + ".h"],
         srcs = [name + ".cc"],
@@ -53,6 +53,9 @@ def tflm_inference_library(
             "//tensorflow/lite/micro:micro_common",
             "//tensorflow/lite/micro:micro_context",
         ],
-        copts = micro_copts(),
+        target_compatible_with = select({
+            "//conditions:default": [],
+            "//:with_compression_enabled": ["@platforms//:incompatible"],
+        }),
         visibility = visibility,
     )
diff --git a/codegen/runtime/BUILD b/codegen/runtime/BUILD
@@ -1,12 +1,11 @@
-load("//tensorflow/lite/micro:build_def.bzl", "micro_copts")
+load("//tensorflow/lite/micro:build_def.bzl", "tflm_cc_library")
 
 package(default_visibility = ["//visibility:public"])
 
-cc_library(
+tflm_cc_library(
     name = "micro_codegen_context",
     srcs = ["micro_codegen_context.cc"],
     hdrs = ["micro_codegen_context.h"],
-    copts = micro_copts(),
     deps = [
         "//tensorflow/lite/c:common",
         "//tensorflow/lite/kernels:op_macros",

diff --git a/python/tflite_micro/BUILD b/python/tflite_micro/BUILD
@@ -7,7 +7,7 @@ load("@rules_python//python:packaging.bzl", "py_package", "py_wheel")
 load("@tflm_pip_deps//:requirements.bzl", "requirement")
 load(
     "//tensorflow/lite/micro:build_def.bzl",
-    "micro_copts",
+    "tflm_cc_library",
 )
 load(
     "//tensorflow:extra_rules.bzl",
@@ -24,15 +24,14 @@ package_group(
     packages = tflm_python_op_resolver_friends(),
 )
 
-cc_library(
+tflm_cc_library(
     name = "python_ops_resolver",
     srcs = [
         "python_ops_resolver.cc",
     ],
     hdrs = [
         "python_ops_resolver.h",
     ],
-    copts = micro_copts(),
     visibility = [
         ":op_resolver_friends",
         "//tensorflow/lite/micro/integration_tests:__subpackages__",

diff --git a/python/tflite_micro/interpreter_wrapper.cc b/python/tflite_micro/interpreter_wrapper.cc
@@ -104,11 +104,6 @@ bool CheckTensor(const TfLiteTensor* tensor) {
     return false;
   }
 
-  if (tensor->sparsity != nullptr) {
-    PyErr_SetString(PyExc_ValueError, "TFLM doesn't support sparse tensors");
-    return false;
-  }
-
   int py_type_num = TfLiteTypeToPyArrayType(tensor->type);
   if (py_type_num == NPY_NOTYPE) {
     PyErr_SetString(PyExc_ValueError, "Unknown tensor type.");

diff --git a/signal/micro/kernels/BUILD b/signal/micro/kernels/BUILD
@@ -1,11 +1,11 @@
 load(
     "//tensorflow/lite/micro:build_def.bzl",
-    "micro_copts",
+    "tflm_cc_library",
 )
 
 package(licenses = ["notice"])
 
-cc_library(
+tflm_cc_library(
     name = "register_signal_ops",
     srcs = [
         "delay.cc",
@@ -31,7 +31,6 @@ cc_library(
         "irfft.h",
         "rfft.h",
     ],
-    copts = micro_copts(),
     visibility = [
         "//tensorflow/lite/micro",
     ],

diff --git a/tensorflow/compiler/mlir/lite/core/api/BUILD b/tensorflow/compiler/mlir/lite/core/api/BUILD
@@ -1,15 +1,19 @@
 load("//tensorflow/lite:build_def.bzl", "tflite_copts")
-load("//tensorflow/lite/micro:build_def.bzl", "micro_copts")
+load(
+    "//tensorflow/lite/micro:build_def.bzl",
+    "tflm_cc_library",
+    "tflm_copts",
+)
 
 package(
     default_visibility = ["//visibility:public"],
     licenses = ["notice"],
 )
 
-cc_library(
+tflm_cc_library(
     name = "error_reporter",
     srcs = ["error_reporter.cc"],
     hdrs = ["error_reporter.h"],
-    copts = tflite_copts() + micro_copts(),
+    copts = tflm_copts() + tflite_copts(),
     deps = [],
 )
diff --git a/tensorflow/lite/core/api/BUILD b/tensorflow/lite/core/api/BUILD
@@ -1,12 +1,16 @@
 load("//tensorflow/lite:build_def.bzl", "tflite_copts")
-load("//tensorflow/lite/micro:build_def.bzl", "micro_copts")
+load(
+    "//tensorflow/lite/micro:build_def.bzl",
+    "tflm_cc_library",
+    "tflm_copts",
+)
 
 package(
     default_visibility = ["//visibility:private"],
     licenses = ["notice"],
 )
 
-cc_library(
+tflm_cc_library(
     name = "api",
     srcs = [
         "flatbuffer_conversions.cc",
@@ -17,7 +21,7 @@ cc_library(
         "flatbuffer_conversions.h",
         "tensor_utils.h",
     ],
-    copts = tflite_copts() + micro_copts(),
+    copts = tflm_copts() + tflite_copts(),
     visibility = ["//visibility:public"],
     deps = [
         ":error_reporter",
@@ -33,13 +37,13 @@ cc_library(
 # also exported by the "api" target, so that targets which only want to depend
 # on these small abstract base class modules can express more fine-grained
 # dependencies without pulling in tensor_utils and flatbuffer_conversions.
-cc_library(
+tflm_cc_library(
     name = "error_reporter",
     hdrs = [
         "error_reporter.h",
         "//tensorflow/compiler/mlir/lite/core/api:error_reporter.h",
     ],
-    copts = tflite_copts() + micro_copts(),
+    copts = tflm_copts() + tflite_copts(),
     visibility = [
         "//visibility:public",
     ],

diff --git a/tensorflow/lite/experimental/microfrontend/lib/BUILD b/tensorflow/lite/experimental/microfrontend/lib/BUILD
@@ -144,7 +144,7 @@ cc_test(
     name = "filterbank_test",
     srcs = ["filterbank_test.cc"],
     # Setting copts for experimental code to [], but this code should be fixed
-    # to build with the default copts (micro_copts())
+    # to build with the default copts
     copts = [],
     deps = [
         ":filterbank",
@@ -156,7 +156,7 @@ cc_test(
     name = "frontend_test",
     srcs = ["frontend_test.cc"],
     # Setting copts for experimental code to [], but this code should be fixed
-    # to build with the default copts (micro_copts())
+    # to build with the default copts
     copts = [],
     deps = [
         ":frontend",
@@ -168,7 +168,7 @@ cc_test(
     name = "log_scale_test",
     srcs = ["log_scale_test.cc"],
     # Setting copts for experimental code to [], but this code should be fixed
-    # to build with the default copts (micro_copts())
+    # to build with the default copts
     copts = [],
     deps = [
         ":log_scale",
@@ -180,7 +180,7 @@ cc_test(
     name = "noise_reduction_test",
     srcs = ["noise_reduction_test.cc"],
     # Setting copts for experimental code to [], but this code should be fixed
-    # to build with the default copts (micro_copts())
+    # to build with the default copts
     copts = [],
     deps = [
         ":noise_reduction",
@@ -192,7 +192,7 @@ cc_test(
     name = "pcan_gain_control_test",
     srcs = ["pcan_gain_control_test.cc"],
     # Setting copts for experimental code to [], but this code should be fixed
-    # to build with the default copts (micro_copts())
+    # to build with the default copts
     copts = [],
     deps = [
         ":pcan_gain_control",
@@ -204,7 +204,7 @@ cc_test(
     name = "window_test",
     srcs = ["window_test.cc"],
     # Setting copts for experimental code to [], but this code should be fixed
-    # to build with the default copts (micro_copts())
+    # to build with the default copts
     copts = [],
     deps = [
         ":window",

diff --git a/tensorflow/lite/kernels/BUILD b/tensorflow/lite/kernels/BUILD
@@ -1,5 +1,9 @@
 load("//tensorflow/lite:build_def.bzl", "tflite_copts")
-load("//tensorflow/lite/micro:build_def.bzl", "micro_copts")
+load(
+    "//tensorflow/lite/micro:build_def.bzl",
+    "tflm_cc_library",
+    "tflm_copts",
+)
 
 package(
     default_visibility = [
@@ -17,15 +21,15 @@ cc_library(
     deps = ["//tensorflow/lite/micro:micro_log"],
 )
 
-cc_library(
+tflm_cc_library(
     name = "kernel_util",
     srcs = [
         "kernel_util.cc",
     ],
     hdrs = [
         "kernel_util.h",
     ],
-    copts = tflite_copts() + micro_copts(),
+    copts = tflm_copts() + tflite_copts(),
     deps = [
         "//tensorflow/lite:array",
         "//tensorflow/lite:kernel_api",

diff --git a/tensorflow/lite/kernels/internal/BUILD b/tensorflow/lite/kernels/internal/BUILD
@@ -1,5 +1,9 @@
 load("//tensorflow/lite:build_def.bzl", "tflite_copts")
-load("//tensorflow/lite/micro:build_def.bzl", "micro_copts")
+load(
+    "//tensorflow/lite/micro:build_def.bzl",
+    "tflm_cc_library",
+    "tflm_copts",
+)
 
 package(
     default_visibility = [
@@ -44,11 +48,11 @@ cc_library(
     copts = tflite_copts(),
 )
 
-cc_library(
+tflm_cc_library(
     name = "quantization_util",
     srcs = ["quantization_util.cc"],
     hdrs = ["quantization_util.h"],
-    copts = tflite_copts() + micro_copts(),
+    copts = tflm_copts() + tflite_copts(),
     deps = [
         ":compatibility",
         ":cppmath",