Add preliminary support for the xcore target.

Author: AdamHillier

larq_compute_engine/mlir/BUILD

@@ -36,28 +36,60 @@ gentbl(
 )
 
 gentbl(
-    name = "prepare_lce_inc_gen",
+    name = "prepare_lce_target_arm_inc_gen",
     tbl_outs = [
-        ("-gen-rewriters", "transforms/generated_prepare.inc"),
+        ("-gen-rewriters", "transforms/generated_prepare_target_arm.inc"),
     ],
     tblgen = "@llvm-project//mlir:mlir-tblgen",
-    td_file = "transforms/prepare_patterns.td",
+    td_file = "transforms/prepare_patterns_target_arm.td",
     td_srcs = [
         "ir/lce_ops.td",
         "transforms/op_removal_patterns.td",
+        "transforms/prepare_patterns_common.td",
         "@llvm-project//mlir:StdOpsTdFiles",
         "@org_tensorflow//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
         "@org_tensorflow//tensorflow/compiler/mlir/lite:tensorflow_lite_ops_td_files",
     ],
 )
 
 gentbl(
-    name = "optimize_lce_inc_gen",
+    name = "prepare_lce_target_other_inc_gen",
     tbl_outs = [
-        ("-gen-rewriters", "transforms/generated_optimize.inc"),
+        ("-gen-rewriters", "transforms/generated_prepare_target_other.inc"),
     ],
     tblgen = "@llvm-project//mlir:mlir-tblgen",
-    td_file = "transforms/optimize_patterns.td",
+    td_file = "transforms/prepare_patterns_common.td",
+    td_srcs = [
+        "ir/lce_ops.td",
+        "transforms/op_removal_patterns.td",
+        "@llvm-project//mlir:StdOpsTdFiles",
+        "@org_tensorflow//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
+        "@org_tensorflow//tensorflow/compiler/mlir/lite:tensorflow_lite_ops_td_files",
+    ],
+)
+
+gentbl(
+    name = "optimize_lce_target_arm_inc_gen",
+    tbl_outs = [
+        ("-gen-rewriters", "transforms/generated_optimize_target_arm.inc"),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "transforms/optimize_patterns_target_arm.td",
+    td_srcs = [
+        "ir/lce_ops.td",
+        "transforms/optimize_patterns_common.td",
+        "@org_tensorflow//tensorflow/compiler/mlir/lite:tensorflow_lite_ops_td_files",
+        "@llvm-project//mlir:StdOpsTdFiles",
+    ],
+)
+
+gentbl(
+    name = "optimize_lce_target_other_inc_gen",
+    tbl_outs = [
+        ("-gen-rewriters", "transforms/generated_optimize_target_other.inc"),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "transforms/optimize_patterns_common.td",
     td_srcs = [
         "ir/lce_ops.td",
         "@org_tensorflow//tensorflow/compiler/mlir/lite:tensorflow_lite_ops_td_files",
@@ -136,7 +168,8 @@ cc_library(
 cc_library(
     name = "larq_compute_engine_prepare",
     srcs = [
-        "transforms/generated_prepare.inc",
+        "transforms/generated_prepare_target_arm.inc",
+        "transforms/generated_prepare_target_other.inc",
         "transforms/prepare_tf.cc",
     ],
     hdrs = [
@@ -157,7 +190,8 @@ cc_library(
 cc_library(
     name = "larq_compute_engine_optimize",
     srcs = [
-        "transforms/generated_optimize.inc",
+        "transforms/generated_optimize_target_arm.inc",
+        "transforms/generated_optimize_target_other.inc",
         "transforms/optimize.cc",
     ],
     hdrs = [

larq_compute_engine/mlir/python/converter.py

@@ -55,6 +55,7 @@ def convert_keras_model(
     *,  # Require remaining arguments to be keyword-only.
     inference_input_type: tf.DType = tf.float32,
     inference_output_type: tf.DType = tf.float32,
+    target: str = "arm",
     experimental_default_int8_range: Optional[Tuple[float, float]] = None,
     experimental_enable_bitpacked_activations: bool = False,
 ) -> bytes:
@@ -73,6 +74,7 @@ def convert_keras_model(
             must be either `tf.float32` or `tf.int8`.
         inference_output_type: Data type of the output layer. Defaults to `tf.float32`,
             must be either `tf.float32` or `tf.int8`.
+        target: Target hardware platform. Must be "arm" or "xcore".
         experimental_default_int8_range: Tuple of integers representing `(min, max)`
             range values for all arrays without a specified range. Intended for
             experimenting with quantization via "dummy quantization". (default None)
@@ -98,6 +100,8 @@ def convert_keras_model(
             "Expected `inference_output_type` to be either `tf.float32` or `tf.int8`, "
             f"got {inference_output_type}."
         )
+    if target not in ("arm", "xcore"):
+        raise ValueError(f'Expected `target` to be "arm" or "xcore", but got {target}.')
 
     if not tf.executing_eagerly():
         raise RuntimeError(
@@ -147,6 +151,7 @@ def convert_keras_model(
         [tensor.shape.as_list() for tensor in input_tensors],
         [get_tensor_name(tensor) for tensor in output_tensors],
         should_quantize,
+        target,
         experimental_default_int8_range,
         experimental_enable_bitpacked_activations,
     )

larq_compute_engine/mlir/python/converter_test.py

@@ -31,6 +31,7 @@ def test_larq_zoo_models(self):
             [[1, 224, 224, 3]],
             ["Identity"],
             False,
+            "arm",
             None,
             False,
         )
@@ -39,6 +40,20 @@ def test_wrong_arg(self):
         with self.assertRaises(ValueError):
             convert_keras_model("./model.h5")
 
+    def test_target_arg(self):
+        with context.eager_mode():
+            model = lqz.sota.QuickNet(weights=None)
+
+            # These should work
+            convert_keras_model(model, target="arm")
+            convert_keras_model(model, target="xcore")
+
+            # Anything else shouldn't
+            with self.assertRaises(
+                ValueError, msg='Expected `target` to be "arm" or "xcore"'
+            ):
+                convert_keras_model(model, target="x86")
+
 
 if __name__ == "__main__":
     unittest.main()

larq_compute_engine/mlir/python/graphdef_tfl_flatbuffer.cc

@@ -41,14 +41,23 @@ pybind11::bytes ConvertGraphDefToTFLiteFlatBuffer(
     const std::vector<string>& input_dtypes,
     const std::vector<std::vector<int>>& input_shapes,
     const std::vector<string>& output_arrays, const bool should_quantize,
-    const pybind11::object& default_ranges,
+    const std::string& target_str, const pybind11::object& default_ranges,
     const bool experimental_enable_bitpacked_activations) {
   GraphDef graphdef;
   if (!tensorflow::LoadProtoFromBuffer(std::string(graphdef_bytes), &graphdef)
            .ok()) {
     throw std::runtime_error("Could not load GraphDef.");
   }
 
+  LCETarget target;
+  if (target_str == "arm") {
+    target = LCETarget::ARM;
+  } else if (target_str == "xcore") {
+    target = LCETarget::XCORE;
+  } else {
+    throw std::runtime_error("Invalid target.");
+  }
+
   GraphImportConfig specs;
   specs.prune_unused_nodes = true;
   specs.convert_legacy_fed_inputs = true;
@@ -88,7 +97,7 @@ pybind11::bytes ConvertGraphDefToTFLiteFlatBuffer(
   }
   mlir::PassManager pm(&context);
   tensorflow::AddTFToLCETFLConversionPasses(
-      quant_specs, &pm, experimental_enable_bitpacked_activations);
+      quant_specs, &pm, target, experimental_enable_bitpacked_activations);
 
   // Convert back to outlined while format for export back to flatbuffer.
   pm.addPass(mlir::TFL::CreateWhileOutlinePass());

larq_compute_engine/mlir/tests/optimize.mlir

@@ -1,4 +1,5 @@
-// RUN: lce-tf-opt %s -tfl-optimize-lce -verify-diagnostics | FileCheck %s
+// RUN: lce-tf-opt %s -tfl-optimize-lce=target-arm=true -verify-diagnostics | FileCheck %s --check-prefixes CHECK,CHECK-ARM
+// RUN: lce-tf-opt %s -tfl-optimize-lce=target-arm=false -verify-diagnostics | FileCheck %s --check-prefixes CHECK,CHECK-OTHER
 
 // CHECK-LABEL: @fuse_add_into_bconv2d
 func @fuse_add_into_bconv2d(%arg0: tensor<256x32x32x1xi32>, %arg1: tensor<16x3x3x3xf32>, %arg2: tensor<16xf32>, %arg3: none) -> tensor<256x30x30x16xf32> {
@@ -13,7 +14,6 @@ func @fuse_add_into_bconv2d(%arg0: tensor<256x32x32x1xi32>, %arg1: tensor<16x3x3
   // CHECK-NEXT: return %0
 }
 
-
 // CHECK-LABEL: @fuse_sub_into_bconv2d
 func @fuse_sub_into_bconv2d(%arg0: tensor<256x32x32x1xi32>, %arg1: tensor<16x3x3x3xf32>, %arg2: tensor<16xf32>, %arg3: none) -> tensor<256x30x30x16xf32> {
   %cst = constant dense<0.5> : tensor<16xf32>
@@ -144,15 +144,19 @@ func @do_not_fuse_relu_into_bconv2d_no_post_activation_multiplier(%arg0: tensor<
   // CHECK-NEXT: return %1
 }
 
-// CHECK-LABEL: @reorder_maxpool_2d_quantize
-func @reorder_maxpool_2d_quantize(%arg0: tensor<256x32x32x65xf32>) -> tensor<256x16x8x3xi32> {
+// CHECK-LABEL: @target_specific_reorder_maxpool_2d_quantize
+func @target_specific_reorder_maxpool_2d_quantize(%arg0: tensor<256x32x32x65xf32>) -> tensor<256x16x8x3xi32> {
   %0 = "tfl.max_pool_2d"(%arg0) {filter_height = 3 : i32, filter_width = 2 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 4 : i32} : (tensor<256x32x32x65xf32>) -> tensor<256x16x8x65xf32>
   %1 = "lq.Quantize"(%0) : (tensor<256x16x8x65xf32>) -> tensor<256x16x8x3xi32>
   return %1 : tensor<256x16x8x3xi32>
 
-  // CHECK-NEXT: %0 = "lq.Quantize"(%arg0) : (tensor<256x32x32x65xf32>) -> tensor<256x32x32x3xi32>
-  // CHECK-NEXT: %1 = "lq.BMaxPool2d"(%0) {filter_height = 3 : i32, filter_width = 2 : i32, padding = "SAME", stride_height = 2 : i32, stride_width = 4 : i32} : (tensor<256x32x32x3xi32>) -> tensor<256x16x8x3xi32>
-  // CHECK-NEXT: return %1
+  // CHECK-ARM-NEXT: %0 = "lq.Quantize"(%arg0) : (tensor<256x32x32x65xf32>) -> tensor<256x32x32x3xi32>
+  // CHECK-ARM-NEXT: %1 = "lq.BMaxPool2d"(%0) {filter_height = 3 : i32, filter_width = 2 : i32, padding = "SAME", stride_height = 2 : i32, stride_width = 4 : i32} : (tensor<256x32x32x3xi32>) -> tensor<256x16x8x3xi32>
+  // CHECK-ARM-NEXT: return %1
+
+  // CHECK-OTHER-NEXT: %0 = "tfl.max_pool_2d"(%arg0) {filter_height = 3 : i32, filter_width = 2 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 4 : i32} : (tensor<256x32x32x65xf32>) -> tensor<256x16x8x65xf32>
+  // CHECK-OTHER-NEXT: %1 = "lq.Quantize"(%0) : (tensor<256x16x8x65xf32>) -> tensor<256x16x8x3xi32>
+  // CHECK-OTHER-NEXT: return %1
 }
 
 // CHECK-LABEL: @do_not_reorder_maxpool_2d_quantize_multiple_uses